diff --git a/ALICE3/DataModel/A3DecayFinderTables.h b/ALICE3/DataModel/A3DecayFinderTables.h index 55229bbb5d4..d1c87c989bd 100644 --- a/ALICE3/DataModel/A3DecayFinderTables.h +++ b/ALICE3/DataModel/A3DecayFinderTables.h @@ -19,9 +19,11 @@ #define ALICE3_DATAMODEL_A3DECAYFINDERTABLES_H_ // O2 includes -#include "Framework/AnalysisDataModel.h" #include "Common/Core/RecoDecay.h" +#include "CommonConstants/PhysicsConstants.h" +#include "Framework/AnalysisDataModel.h" + enum a3selectionBit : uint32_t { kDCAxy = 0, kInnerTOFPion, kInnerTOFKaon, @@ -51,6 +53,187 @@ enum a3selectionBit : uint32_t { kDCAxy = 0, namespace o2::aod { + +// general decay properties +namespace a3_hf_cand +{ +// collision properties +DECLARE_SOA_INDEX_COLUMN(Collision, collision); //! +// secondary vertex +DECLARE_SOA_COLUMN(XSecondaryVertex, xSecondaryVertex, double); //! +DECLARE_SOA_COLUMN(YSecondaryVertex, ySecondaryVertex, double); //! +DECLARE_SOA_COLUMN(ZSecondaryVertex, zSecondaryVertex, double); //! +DECLARE_SOA_DYNAMIC_COLUMN(RSecondaryVertex, rSecondaryVertex, //! + [](float xVtxS, float yVtxS) -> float { return RecoDecay::sqrtSumOfSquares(xVtxS, yVtxS); }); +DECLARE_SOA_COLUMN(Chi2PCA, chi2PCA, float); //! sum of (non-weighted) distances of the secondary vertex to its prongs +// prong properties +DECLARE_SOA_COLUMN(PxProng0, pxProng0, float); //! +DECLARE_SOA_COLUMN(PyProng0, pyProng0, float); //! +DECLARE_SOA_COLUMN(PzProng0, pzProng0, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(PtProng0, ptProng0, //! + [](float px, float py) -> float { return RecoDecay::pt(px, py); }); +DECLARE_SOA_DYNAMIC_COLUMN(Pt2Prong0, pt2Prong0, //! + [](float px, float py) -> float { return RecoDecay::pt2(px, py); }); +DECLARE_SOA_DYNAMIC_COLUMN(PVectorProng0, pVectorProng0, //! + [](float px, float py, float pz) -> std::array { return std::array{px, py, pz}; }); +DECLARE_SOA_COLUMN(ImpactParameterY0, impactParameterY0, float); //! +DECLARE_SOA_COLUMN(ErrorImpactParameterY0, errorImpactParameterY0, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(ImpactParameterNormalised0, impactParameterNormalised0, //! + [](float dca, float err) -> float { return dca / err; }); +DECLARE_SOA_COLUMN(ImpactParameterZ0, impactParameterZ0, float); //! +DECLARE_SOA_COLUMN(ErrorImpactParameterZ0, errorImpactParameterZ0, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(ImpactParameterZNormalised0, impactParameterZNormalised0, //! + [](float dca, float err) -> float { return dca / err; }); +DECLARE_SOA_COLUMN(PxProng1, pxProng1, float); //! +DECLARE_SOA_COLUMN(PyProng1, pyProng1, float); //! +DECLARE_SOA_COLUMN(PzProng1, pzProng1, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(PtProng1, ptProng1, //! + [](float px, float py) -> float { return RecoDecay::pt(px, py); }); +DECLARE_SOA_DYNAMIC_COLUMN(Pt2Prong1, pt2Prong1, //! + [](float px, float py) -> float { return RecoDecay::pt2(px, py); }); +DECLARE_SOA_DYNAMIC_COLUMN(PVectorProng1, pVectorProng1, //! + [](float px, float py, float pz) -> std::array { return std::array{px, py, pz}; }); +DECLARE_SOA_COLUMN(ImpactParameterY1, impactParameterY1, float); //! +DECLARE_SOA_COLUMN(ErrorImpactParameterY1, errorImpactParameterY1, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(ImpactParameterNormalised1, impactParameterNormalised1, //! + [](float dca, float err) -> float { return dca / err; }); +DECLARE_SOA_COLUMN(ImpactParameterZ1, impactParameterZ1, float); //! +DECLARE_SOA_COLUMN(ErrorImpactParameterZ1, errorImpactParameterZ1, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(ImpactParameterZNormalised1, impactParameterZNormalised1, //! + [](float dca, float err) -> float { return dca / err; }); +DECLARE_SOA_COLUMN(PxProng2, pxProng2, float); //! +DECLARE_SOA_COLUMN(PyProng2, pyProng2, float); //! +DECLARE_SOA_COLUMN(PzProng2, pzProng2, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(PtProng2, ptProng2, //! + [](float px, float py) -> float { return RecoDecay::pt(px, py); }); +DECLARE_SOA_DYNAMIC_COLUMN(Pt2Prong2, pt2Prong2, //! + [](float px, float py) -> float { return RecoDecay::pt2(px, py); }); +DECLARE_SOA_DYNAMIC_COLUMN(PVectorProng2, pVectorProng2, //! + [](float px, float py, float pz) -> std::array { return std::array{px, py, pz}; }); +DECLARE_SOA_COLUMN(ImpactParameterY2, impactParameterY2, float); //! +DECLARE_SOA_COLUMN(ErrorImpactParameterY2, errorImpactParameterY2, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(ImpactParameterNormalised2, impactParameterNormalised2, //! + [](float dca, float err) -> float { return dca / err; }); +DECLARE_SOA_COLUMN(ImpactParameterZ2, impactParameterZ2, float); //! +DECLARE_SOA_COLUMN(ErrorImpactParameterZ2, errorImpactParameterZ2, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(ImpactParameterZNormalised2, impactParameterZNormalised2, //! + [](float dca, float err) -> float { return dca / err; }); + +/// prong PID nsigma +DECLARE_SOA_COLUMN(NSigTrkPi0, nSigTrkPi0, float); //! +DECLARE_SOA_COLUMN(NSigTrkKa0, nSigTrkKa0, float); //! +DECLARE_SOA_COLUMN(NSigTrkPr0, nSigTrkPr0, float); //! +DECLARE_SOA_COLUMN(NSigTrkPi1, nSigTrkPi1, float); //! +DECLARE_SOA_COLUMN(NSigTrkKa1, nSigTrkKa1, float); //! +DECLARE_SOA_COLUMN(NSigTrkPr1, nSigTrkPr1, float); //! +DECLARE_SOA_COLUMN(NSigTrkPi2, nSigTrkPi2, float); //! +DECLARE_SOA_COLUMN(NSigTrkKa2, nSigTrkKa2, float); //! +DECLARE_SOA_COLUMN(NSigTrkPr2, nSigTrkPr2, float); //! +DECLARE_SOA_COLUMN(NSigRichPi0, nSigRichPi0, float); //! +DECLARE_SOA_COLUMN(NSigRichKa0, nSigRichKa0, float); //! +DECLARE_SOA_COLUMN(NSigRichPr0, nSigRichPr0, float); //! +DECLARE_SOA_COLUMN(NSigRichPi1, nSigRichPi1, float); //! +DECLARE_SOA_COLUMN(NSigRichKa1, nSigRichKa1, float); //! +DECLARE_SOA_COLUMN(NSigRichPr1, nSigRichPr1, float); //! +DECLARE_SOA_COLUMN(NSigRichPi2, nSigRichPi2, float); //! +DECLARE_SOA_COLUMN(NSigRichKa2, nSigRichKa2, float); //! +DECLARE_SOA_COLUMN(NSigRichPr2, nSigRichPr2, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPi0, nSigInnTofPi0, float); //! +DECLARE_SOA_COLUMN(NSigInnTofKa0, nSigInnTofKa0, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPr0, nSigInnTofPr0, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPi1, nSigInnTofPi1, float); //! +DECLARE_SOA_COLUMN(NSigInnTofKa1, nSigInnTofKa1, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPr1, nSigInnTofPr1, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPi2, nSigInnTofPi2, float); //! +DECLARE_SOA_COLUMN(NSigInnTofKa2, nSigInnTofKa2, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPr2, nSigInnTofPr2, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPi0, nSigOutTofPi0, float); //! +DECLARE_SOA_COLUMN(NSigOutTofKa0, nSigOutTofKa0, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPr0, nSigOutTofPr0, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPi1, nSigOutTofPi1, float); //! +DECLARE_SOA_COLUMN(NSigOutTofKa1, nSigOutTofKa1, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPr1, nSigOutTofPr1, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPi2, nSigOutTofPi2, float); //! +DECLARE_SOA_COLUMN(NSigOutTofKa2, nSigOutTofKa2, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPr2, nSigOutTofPr2, float); //! + +// candidate properties +DECLARE_SOA_DYNAMIC_COLUMN(Pt, pt, //! + [](float px, float py) -> float { return RecoDecay::pt(px, py); }); +DECLARE_SOA_DYNAMIC_COLUMN(Pt2, pt2, //! + [](float px, float py) -> float { return RecoDecay::pt2(px, py); }); +DECLARE_SOA_DYNAMIC_COLUMN(P, p, //! + [](float px, float py, float pz) -> float { return RecoDecay::p(px, py, pz); }); +DECLARE_SOA_DYNAMIC_COLUMN(P2, p2, //! + [](float px, float py, float pz) -> float { return RecoDecay::p2(px, py, pz); }); +DECLARE_SOA_DYNAMIC_COLUMN(PVector, pVector, //! + [](float px, float py, float pz) -> std::array { return std::array{px, py, pz}; }); +DECLARE_SOA_DYNAMIC_COLUMN(Eta, eta, //! + [](float px, float py, float pz) -> float { return RecoDecay::eta(std::array{px, py, pz}); }); +DECLARE_SOA_DYNAMIC_COLUMN(Phi, phi, //! + [](float px, float py) -> float { return RecoDecay::phi(px, py); }); +DECLARE_SOA_DYNAMIC_COLUMN(Y, y, //! + [](float px, float py, float pz, double m) -> float { return RecoDecay::y(std::array{px, py, pz}, m); }); +DECLARE_SOA_DYNAMIC_COLUMN(E, e, //! + [](float px, float py, float pz, double m) -> float { return RecoDecay::e(px, py, pz, m); }); +DECLARE_SOA_DYNAMIC_COLUMN(E2, e2, //! + [](float px, float py, float pz, double m) -> float { return RecoDecay::e2(px, py, pz, m); }); +DECLARE_SOA_DYNAMIC_COLUMN(DecayLength, decayLength, //! + [](float xVtxP, float yVtxP, float zVtxP, float xVtxS, float yVtxS, float zVtxS) -> float { return RecoDecay::distance(std::array{xVtxP, yVtxP, zVtxP}, std::array{xVtxS, yVtxS, zVtxS}); }); +DECLARE_SOA_DYNAMIC_COLUMN(DecayLengthXY, decayLengthXY, //! + [](float xVtxP, float yVtxP, float xVtxS, float yVtxS) -> float { return RecoDecay::distanceXY(std::array{xVtxP, yVtxP}, std::array{xVtxS, yVtxS}); }); +DECLARE_SOA_DYNAMIC_COLUMN(DecayLengthNormalised, decayLengthNormalised, //! + [](float xVtxP, float yVtxP, float zVtxP, float xVtxS, float yVtxS, float zVtxS, float err) -> float { return RecoDecay::distance(std::array{xVtxP, yVtxP, zVtxP}, std::array{xVtxS, yVtxS, zVtxS}) / err; }); +DECLARE_SOA_DYNAMIC_COLUMN(DecayLengthXYNormalised, decayLengthXYNormalised, //! + [](float xVtxP, float yVtxP, float xVtxS, float yVtxS, float err) -> float { return RecoDecay::distanceXY(std::array{xVtxP, yVtxP}, std::array{xVtxS, yVtxS}) / err; }); +DECLARE_SOA_COLUMN(ErrorDecayLength, errorDecayLength, float); //! +DECLARE_SOA_COLUMN(ErrorDecayLengthXY, errorDecayLengthXY, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(Cpa, cpa, //! + [](float xVtxP, float yVtxP, float zVtxP, float xVtxS, float yVtxS, float zVtxS, float px, float py, float pz) -> float { return RecoDecay::cpa(std::array{xVtxP, yVtxP, zVtxP}, std::array{xVtxS, yVtxS, zVtxS}, std::array{px, py, pz}); }); +DECLARE_SOA_DYNAMIC_COLUMN(CpaXY, cpaXY, //! + [](float xVtxP, float yVtxP, float xVtxS, float yVtxS, float px, float py) -> float { return RecoDecay::cpaXY(std::array{xVtxP, yVtxP}, std::array{xVtxS, yVtxS}, std::array{px, py}); }); +DECLARE_SOA_DYNAMIC_COLUMN(Ct, ct, //! + [](float xVtxP, float yVtxP, float zVtxP, float xVtxS, float yVtxS, float zVtxS, float px, float py, float pz, double m) -> float { return RecoDecay::ct(std::array{px, py, pz}, RecoDecay::distance(std::array{xVtxP, yVtxP, zVtxP}, std::array{xVtxS, yVtxS, zVtxS}), m); }); +DECLARE_SOA_DYNAMIC_COLUMN(ImpactParameterXY, impactParameterXY, //! + [](float xVtxP, float yVtxP, float zVtxP, float xVtxS, float yVtxS, float zVtxS, float px, float py, float pz) -> float { return RecoDecay::impParXY(std::array{xVtxP, yVtxP, zVtxP}, std::array{xVtxS, yVtxS, zVtxS}, std::array{px, py, pz}); }); + +// ML scores columns +DECLARE_SOA_COLUMN(MlScore0, mlScore0, float); //! +DECLARE_SOA_COLUMN(MlScore1, mlScore1, float); //! +DECLARE_SOA_COLUMN(MlScore2, mlScore2, float); //! +} // namespace a3_hf_cand + +#define HFCAND_COLUMNS \ + a3_hf_cand::CollisionId, \ + collision::PosX, collision::PosY, collision::PosZ, \ + a3_hf_cand::XSecondaryVertex, a3_hf_cand::YSecondaryVertex, a3_hf_cand::ZSecondaryVertex, \ + a3_hf_cand::ErrorDecayLength, a3_hf_cand::ErrorDecayLengthXY, \ + a3_hf_cand::Chi2PCA, \ + /* dynamic columns */ a3_hf_cand::RSecondaryVertex, \ + a3_hf_cand::DecayLength, \ + a3_hf_cand::DecayLengthXY, \ + a3_hf_cand::DecayLengthNormalised, \ + a3_hf_cand::DecayLengthXYNormalised + +// general columns +#define HFPRONG0_COLUMNS \ + a3_hf_cand::ImpactParameterNormalised0, \ + a3_hf_cand::PtProng0, \ + a3_hf_cand::Pt2Prong0, \ + a3_hf_cand::PVectorProng0 + +#define HFPRONG1_COLUMNS \ + a3_hf_cand::ImpactParameterNormalised1, \ + a3_hf_cand::PtProng1, \ + a3_hf_cand::Pt2Prong1, \ + a3_hf_cand::PVectorProng1 + +#define HFPRONG2_COLUMNS \ + a3_hf_cand::ImpactParameterNormalised2, \ + a3_hf_cand::PtProng2, \ + a3_hf_cand::Pt2Prong2, \ + a3_hf_cand::PVectorProng2 + namespace a3DecayMap { DECLARE_SOA_COLUMN(DecayMap, decayMap, uint32_t); //! simple map to process passing / not passing criteria @@ -118,6 +301,127 @@ DECLARE_SOA_COLUMN(McTruthInfo, mcTruthInfo, int); //! 0 for bkg, 1 for true D0, DECLARE_SOA_TABLE(Alice3D0MCTruth, "AOD", "ALICE3D0MCTRUTH", //! a3D0MCTruth::McTruthInfo); //! +// Now let's define the Lc to pKpi table +namespace a3_hf_cand_3prong +{ +DECLARE_SOA_INDEX_COLUMN(Collision, collision); //! +DECLARE_SOA_COLUMN(Px, px, float); //! +DECLARE_SOA_COLUMN(Py, py, float); //! +DECLARE_SOA_COLUMN(Pz, pz, float); //! +DECLARE_SOA_COLUMN(Pt, pt, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(M, m, + [](float px0, float py0, float pz0, + float px1, float py1, float pz1, + float px2, float py2, float pz2, + const std::array& m) -> float { return RecoDecay::m(std::array{std::array{px0, py0, pz0}, + std::array{px1, py1, pz1}, + std::array{px2, py2, pz2}}, + m); }); +DECLARE_SOA_DYNAMIC_COLUMN(E, e, //! + [](float px, float py, float pz, const float m) -> float { return RecoDecay::e(px, py, pz, m); }); +DECLARE_SOA_COLUMN(Eta, eta, float); //! +DECLARE_SOA_COLUMN(Phi, phi, float); //! +DECLARE_SOA_DYNAMIC_COLUMN(Y, y, + [](float px, float py, float pz, const float m) -> float { return RecoDecay::y(std::array{px, py, pz}, m); }); +} // namespace a3_hf_cand_3prong +DECLARE_SOA_TABLE(Alice3Cand3Ps, "AOD", "ALICE3CAND3P", //! + o2::soa::Index<>, + // general candidate properties + HFCAND_COLUMNS, + HFPRONG0_COLUMNS, + HFPRONG1_COLUMNS, + HFPRONG2_COLUMNS, + // candidate kinematics + a3_hf_cand_3prong::Eta, + a3_hf_cand_3prong::Phi, + a3_hf_cand_3prong::Pt, + // prong properties + a3_hf_cand::PxProng0, a3_hf_cand::PyProng0, a3_hf_cand::PzProng0, // proton track + a3_hf_cand::PxProng1, a3_hf_cand::PyProng1, a3_hf_cand::PzProng1, // kaon track + a3_hf_cand::PxProng2, a3_hf_cand::PyProng2, a3_hf_cand::PzProng2, // pion track + a3_hf_cand::ImpactParameterY0, a3_hf_cand::ImpactParameterY1, a3_hf_cand::ImpactParameterY2, + a3_hf_cand::ErrorImpactParameterY0, a3_hf_cand::ErrorImpactParameterY1, a3_hf_cand::ErrorImpactParameterY2, + a3_hf_cand::ImpactParameterZ0, a3_hf_cand::ImpactParameterZ1, a3_hf_cand::ImpactParameterZ2, + a3_hf_cand::ErrorImpactParameterZ0, a3_hf_cand::ErrorImpactParameterZ1, a3_hf_cand::ErrorImpactParameterZ2, + // Candidate momenta + a3_hf_cand_3prong::Px, a3_hf_cand_3prong::Py, a3_hf_cand_3prong::Pz, + // dynamic candidate properties + a3_hf_cand::Cpa, + a3_hf_cand::CpaXY, + a3_hf_cand::ImpactParameterXY, + a3_hf_cand_3prong::Y, + a3_hf_cand_3prong::M, + a3_hf_cand_3prong::E); + +namespace a3_hf_sel_3prong +{ +DECLARE_SOA_COLUMN(IsSelMassHypo0, isSelMassHypo0, bool); //! +DECLARE_SOA_COLUMN(IsSelMassHypo1, isSelMassHypo1, bool); //! + +// PID selection +enum PidSels { + None = 0, + TrkProng0, + RichProng0, + InnTofProng0, + OutTofProng0, + TrkProng1, + RichProng1, + InnTofProng1, + OutTofProng1, + TrkProng2, + RichProng2, + InnTofProng2, + OutTofProng2, + NPidSelections +}; +DECLARE_SOA_COLUMN(PidBitMask, pidBitMask, uint32_t); //! +} // namespace a3_hf_sel_3prong +DECLARE_SOA_TABLE(Alice3Sel3Ps, "AOD", "ALICE3SEL3P", //! + a3_hf_sel_3prong::IsSelMassHypo0, + a3_hf_sel_3prong::IsSelMassHypo1, + a3_hf_sel_3prong::PidBitMask); + +namespace a3_mc_truth +{ +DECLARE_SOA_COLUMN(OriginMcRec, originMcRec, int); //! +DECLARE_SOA_COLUMN(BHadMotherPtRec, bHadMotherPtRec, float); //! +DECLARE_SOA_COLUMN(FlagMcRec, flagMcRec, int); //! +DECLARE_SOA_COLUMN(OriginMcGen, originMcGen, int); //! +DECLARE_SOA_COLUMN(BHadMotherPtGen, bHadMotherPtGen, float); //! +DECLARE_SOA_COLUMN(FlagMcGen, flagMcGen, int); //! +} // namespace a3_mc_truth +DECLARE_SOA_TABLE(Alice3McRecFlags, "AOD", "ALICE3MCRECFLAG", //! + a3_mc_truth::OriginMcRec, + a3_mc_truth::BHadMotherPtRec, + a3_mc_truth::FlagMcRec); + +DECLARE_SOA_TABLE(Alice3McGenFlags, "AOD", "ALICE3MCGENFLAG", //! + a3_mc_truth::OriginMcGen, + a3_mc_truth::BHadMotherPtGen, + a3_mc_truth::FlagMcGen); + +DECLARE_SOA_TABLE(Alice3PidLcs, "AOD", "ALICE3PIDLC", //! + a3_hf_cand::NSigTrkPr0, + a3_hf_cand::NSigRichPr0, + a3_hf_cand::NSigInnTofPr0, + a3_hf_cand::NSigOutTofPr0, + a3_hf_cand::NSigTrkKa1, + a3_hf_cand::NSigRichKa1, + a3_hf_cand::NSigInnTofKa1, + a3_hf_cand::NSigOutTofKa1, + a3_hf_cand::NSigTrkPi2, + a3_hf_cand::NSigRichPi2, + a3_hf_cand::NSigInnTofPi2, + a3_hf_cand::NSigOutTofPi2); + +DECLARE_SOA_TABLE(Alice3Ml3Ps, "AOD", "ALICE3ML3P", //! + a3_hf_cand::MlScore0, + a3_hf_cand::MlScore1, + a3_hf_cand::MlScore2); + } // namespace o2::aod #endif // ALICE3_DATAMODEL_A3DECAYFINDERTABLES_H_ diff --git a/ALICE3/ML/HfMlResponse3Prong.h b/ALICE3/ML/HfMlResponse3Prong.h new file mode 100644 index 00000000000..24c7b479213 --- /dev/null +++ b/ALICE3/ML/HfMlResponse3Prong.h @@ -0,0 +1,222 @@ +// 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 Alice3MlRResponse3Prong.h +/// \brief Class to compute the ML response for HF 3-prong candidates +/// \author Marcello Di Costanzo , Polytechnic University of Turin and INFN Turin + +#ifndef ALICE3_ML_HFMLRESPONSE3PRONG_H_ +#define ALICE3_ML_HFMLRESPONSE3PRONG_H_ + +#include "Tools/ML/MlResponse.h" + +#include +#include +#include +#include + +// Fill the map of available input features +// the key is the feature's name (std::string) +// the value is the corresponding value in EnumInputFeatures +#define FILL_MAP_3PRONG(FEATURE) \ + { \ + #FEATURE, static_cast(InputFeatures3Prong::FEATURE) \ + } + +// Specific case of CHECK_AND_FILL_ML_ALICE3_FULL(OBJECT, FEATURE, GETTER) +// where OBJECT is named candidate and FEATURE = GETTER +#define CHECK_AND_FILL_ML_ALICE3(GETTER) \ + case static_cast(InputFeatures3Prong::GETTER): { \ + inputFeatures.emplace_back(candidate.GETTER()); \ + break; \ + } + +namespace o2::analysis +{ +enum class InputFeatures3Prong : uint8_t { + ptProng0 = 0, + ptProng1, + ptProng2, + impactParameterY0, + impactParameterY1, + impactParameterY2, + impactParameterZ0, + impactParameterZ1, + impactParameterZ2, + decayLength, + decayLengthXY, + decayLengthXYNormalised, + cpa, + cpaXY, + chi2PCA, + nSigRichPr0, // 0 + nSigRichKa0, // 0 + nSigRichPi0, // 0 + nSigRichPr1, // 1 + nSigRichKa1, // 1 + nSigRichPi1, // 1 + nSigRichPr2, // 2 + nSigRichKa2, // 2 + nSigRichPi2, // 2 + nSigInnTofPr0, // 0 + nSigInnTofKa0, // 0 + nSigInnTofPi0, // 0 + nSigInnTofPr1, // 1 + nSigInnTofKa1, // 1 + nSigInnTofPi1, // 1 + nSigInnTofPr2, // 2 + nSigInnTofKa2, // 2 + nSigInnTofPi2, // 2 + nSigOutTofPr0, // 0 + nSigOutTofKa0, // 0 + nSigOutTofPi0, // 0 + nSigOutTofPr1, // 1 + nSigOutTofKa1, // 1 + nSigOutTofPi1, // 1 + nSigOutTofPr2, // 2 + nSigOutTofKa2, // 2 + nSigOutTofPi2, // 2 + nSigTrkPr0, // 0 + nSigTrkKa0, // 0 + nSigTrkPi0, // 0 + nSigTrkPr1, // 1 + nSigTrkKa1, // 1 + nSigTrkPi1, // 1 + nSigTrkPr2, // 2 + nSigTrkKa2, // 2 + nSigTrkPi2 // 2 +}; + +template +class HfMlResponse3Prong : public MlResponse +{ + public: + /// Default constructor + HfMlResponse3Prong() = default; + /// Default destructor + virtual ~HfMlResponse3Prong() = default; + + /// Method to get the input features vector needed for ML inference + /// \tparam T1 type of the 3-prong candidate + /// \param candidate is the 3-prong candidate + /// \return inputFeatures vector + template + std::vector getInputFeatures(T1 const& candidate) + { + std::vector inputFeatures; + + for (const auto& idx : MlResponse::mCachedIndices) { + switch (idx) { + CHECK_AND_FILL_ML_ALICE3(ptProng0); + CHECK_AND_FILL_ML_ALICE3(ptProng1); + CHECK_AND_FILL_ML_ALICE3(ptProng2); + CHECK_AND_FILL_ML_ALICE3(impactParameterY0); + CHECK_AND_FILL_ML_ALICE3(impactParameterY1); + CHECK_AND_FILL_ML_ALICE3(impactParameterY2); + CHECK_AND_FILL_ML_ALICE3(impactParameterZ0); + CHECK_AND_FILL_ML_ALICE3(impactParameterZ1); + CHECK_AND_FILL_ML_ALICE3(impactParameterZ2); + CHECK_AND_FILL_ML_ALICE3(decayLength); + CHECK_AND_FILL_ML_ALICE3(decayLengthXY); + CHECK_AND_FILL_ML_ALICE3(decayLengthXYNormalised); + CHECK_AND_FILL_ML_ALICE3(cpa); + CHECK_AND_FILL_ML_ALICE3(cpaXY); + CHECK_AND_FILL_ML_ALICE3(chi2PCA); + // TRACKER PID variables + CHECK_AND_FILL_ML_ALICE3(nSigTrkPr0); + CHECK_AND_FILL_ML_ALICE3(nSigTrkKa1); + CHECK_AND_FILL_ML_ALICE3(nSigTrkPi2); + // RICH PID variables + CHECK_AND_FILL_ML_ALICE3(nSigRichPr0); + CHECK_AND_FILL_ML_ALICE3(nSigRichKa1); + CHECK_AND_FILL_ML_ALICE3(nSigRichPi2); + // INNER TOF PID variables + CHECK_AND_FILL_ML_ALICE3(nSigInnTofPr0); + CHECK_AND_FILL_ML_ALICE3(nSigInnTofKa1); + CHECK_AND_FILL_ML_ALICE3(nSigInnTofPi2); + // OUTER TOF PID variables + CHECK_AND_FILL_ML_ALICE3(nSigOutTofPr0); + CHECK_AND_FILL_ML_ALICE3(nSigOutTofKa1); + CHECK_AND_FILL_ML_ALICE3(nSigOutTofPi2); + } + } + return inputFeatures; + } + + protected: + /// Method to fill the map of available input features + void setAvailableInputFeatures() + { + MlResponse::mAvailableInputFeatures = { + FILL_MAP_3PRONG(ptProng0), + FILL_MAP_3PRONG(ptProng1), + FILL_MAP_3PRONG(ptProng2), + FILL_MAP_3PRONG(impactParameterY0), + FILL_MAP_3PRONG(impactParameterY1), + FILL_MAP_3PRONG(impactParameterY2), + FILL_MAP_3PRONG(impactParameterZ0), + FILL_MAP_3PRONG(impactParameterZ1), + FILL_MAP_3PRONG(impactParameterZ2), + FILL_MAP_3PRONG(decayLength), + FILL_MAP_3PRONG(decayLengthXY), + FILL_MAP_3PRONG(decayLengthXYNormalised), + FILL_MAP_3PRONG(cpa), + FILL_MAP_3PRONG(cpaXY), + FILL_MAP_3PRONG(chi2PCA), + // RICH PID variables + FILL_MAP_3PRONG(nSigRichPr0), + FILL_MAP_3PRONG(nSigRichKa0), + FILL_MAP_3PRONG(nSigRichPi0), + FILL_MAP_3PRONG(nSigRichPr1), + FILL_MAP_3PRONG(nSigRichKa1), + FILL_MAP_3PRONG(nSigRichPi1), + FILL_MAP_3PRONG(nSigRichPr2), + FILL_MAP_3PRONG(nSigRichKa2), + FILL_MAP_3PRONG(nSigRichPi2), + // INNER TOF PID variables + FILL_MAP_3PRONG(nSigInnTofPr0), + FILL_MAP_3PRONG(nSigInnTofKa0), + FILL_MAP_3PRONG(nSigInnTofPi0), + FILL_MAP_3PRONG(nSigInnTofPr1), + FILL_MAP_3PRONG(nSigInnTofKa1), + FILL_MAP_3PRONG(nSigInnTofPi1), + FILL_MAP_3PRONG(nSigInnTofPr2), + FILL_MAP_3PRONG(nSigInnTofKa2), + FILL_MAP_3PRONG(nSigInnTofPi2), + // OUTER TOF PID variables + FILL_MAP_3PRONG(nSigOutTofPr0), + FILL_MAP_3PRONG(nSigOutTofKa0), + FILL_MAP_3PRONG(nSigOutTofPi0), + FILL_MAP_3PRONG(nSigOutTofPr1), + FILL_MAP_3PRONG(nSigOutTofKa1), + FILL_MAP_3PRONG(nSigOutTofPi1), + FILL_MAP_3PRONG(nSigOutTofPr2), + FILL_MAP_3PRONG(nSigOutTofKa2), + FILL_MAP_3PRONG(nSigOutTofPi2), + // TRACKER PID variables + FILL_MAP_3PRONG(nSigTrkPr0), + FILL_MAP_3PRONG(nSigTrkKa0), + FILL_MAP_3PRONG(nSigTrkPi0), + FILL_MAP_3PRONG(nSigTrkPr1), + FILL_MAP_3PRONG(nSigTrkKa1), + FILL_MAP_3PRONG(nSigTrkPi1), + FILL_MAP_3PRONG(nSigTrkPr2), + FILL_MAP_3PRONG(nSigTrkKa2), + FILL_MAP_3PRONG(nSigTrkPi2)}; + } +}; + +} // namespace o2::analysis + +#undef FILL_MAP_3PRONG +#undef CHECK_AND_FILL_ML_ALICE3 + +#endif // ALICE3_ML_HFMLRESPONSE3PRONG_H_ diff --git a/ALICE3/TableProducer/CMakeLists.txt b/ALICE3/TableProducer/CMakeLists.txt index 4d2c473fe3a..61f46fba7cd 100644 --- a/ALICE3/TableProducer/CMakeLists.txt +++ b/ALICE3/TableProducer/CMakeLists.txt @@ -46,7 +46,17 @@ o2physics_add_dpl_workflow(alice3-correlatorddbar PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DCAFitter COMPONENT_NAME Analysis) +o2physics_add_dpl_workflow(alice3-hf-selector-3prong + SOURCES alice3HfSelector3Prong.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DCAFitter O2Physics::MLCore + COMPONENT_NAME Analysis) + +o2physics_add_dpl_workflow(alice3-hf-tree-creator-3prong + SOURCES alice3HfTreeCreator3Prong.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DCAFitter + COMPONENT_NAME Analysis) + o2physics_add_dpl_workflow(alice3-tracking-translator SOURCES alice3TrackingTranslator.cxx - PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DCAFitter + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore COMPONENT_NAME Analysis) diff --git a/ALICE3/TableProducer/alice3-decayfinder.cxx b/ALICE3/TableProducer/alice3-decayfinder.cxx index 2d12ea9ed14..a0a6e319055 100644 --- a/ALICE3/TableProducer/alice3-decayfinder.cxx +++ b/ALICE3/TableProducer/alice3-decayfinder.cxx @@ -17,37 +17,42 @@ // HF decays. Work in progress: use at your own risk! // -#include -#include -#include -#include -#include -#include - -#include "Framework/runDataProcessing.h" -#include "Framework/RunningWorkflowInfo.h" -#include "Framework/AnalysisTask.h" -#include "Framework/AnalysisDataModel.h" -#include "Framework/ASoAHelpers.h" -#include "DCAFitter/DCAFitterN.h" -#include "ReconstructionDataFormats/Track.h" +#include "ALICE3/DataModel/A3DecayFinderTables.h" +#include "ALICE3/DataModel/OTFPIDTrk.h" +#include "ALICE3/DataModel/OTFRICH.h" +#include "ALICE3/DataModel/OTFTOF.h" +#include "ALICE3/DataModel/RICH.h" +#include "ALICE3/Utils/utilsHfAlice3.h" #include "Common/Core/RecoDecay.h" -#include "Common/Core/trackUtilities.h" -#include "PWGLF/DataModel/LFStrangenessTables.h" -#include "PWGLF/DataModel/LFParticleIdentification.h" #include "Common/Core/TrackSelection.h" +#include "Common/Core/trackUtilities.h" #include "Common/DataModel/TrackSelectionTables.h" -#include "DetectorsBase/Propagator.h" -#include "DetectorsBase/GeometryManager.h" -#include "DataFormatsParameters/GRPObject.h" -#include "DataFormatsParameters/GRPMagField.h" + #include "CCDB/BasicCCDBManager.h" +#include "DCAFitter/DCAFitterN.h" #include "DataFormatsCalibration/MeanVertexObject.h" -#include "ALICE3/DataModel/OTFTOF.h" -#include "ALICE3/DataModel/RICH.h" -#include "ALICE3/DataModel/A3DecayFinderTables.h" +#include "DataFormatsParameters/GRPMagField.h" +#include "DataFormatsParameters/GRPObject.h" +#include "DetectorsBase/GeometryManager.h" +#include "DetectorsBase/Propagator.h" +#include "Framework/ASoAHelpers.h" +#include "Framework/AnalysisDataModel.h" +#include "Framework/AnalysisTask.h" +#include "Framework/RunningWorkflowInfo.h" +#include "Framework/runDataProcessing.h" +#include "ReconstructionDataFormats/Track.h" + +#include +#include +#include +#include +#include +#include +#include +#include using namespace o2; +using namespace o2::analysis; using namespace o2::framework; using namespace o2::constants::physics; using namespace o2::framework::expressions; @@ -58,13 +63,8 @@ using std::array; #define bitoff(var, nbit) ((var) &= ~(static_cast(1) << (nbit))) //((a) &= ~(1ULL<<(b))) // #define bitcheck(var, nbit) ((var) & (static_cast(1) << (nbit))) -using FullTracksExt = soa::Join; - // For MC association in pre-selection -using labeledTracks = soa::Join; -using tofTracks = soa::Join; -using richTracks = soa::Join; -using alice3tracks = soa::Join; +using Alice3TracksWPid = soa::Join; struct alice3decayFinder { SliceCache cache; @@ -72,11 +72,24 @@ struct alice3decayFinder { Produces candidateD0meson; // contains D0 and D0bar selected candidates (separated, i.e. each row with a single mass hypothesis) Produces selectionOutcome; // flags for isSelD0 and isSelD0bar Produces mcTruthOutcome; // contains MC truth info (is true D0, true D0bar, or bkg) - + Produces candidate3Prong; // contains Lc selected candidates + Produces mcRecFlags; // contains MC truth info (is true Lc, or bkg) + Produces pidInfoLcDaugs; // contains PID info for Lc candidates + Produces mcGenFlags; // contains MC gen info for 3-prong candidates + + // Vertexing + Configurable propagateToPCA{"propagateToPCA", true, "create tracks version propagated to PCA"}; + Configurable useAbsDCA{"useAbsDCA", false, "Minimise abs. distance rather than chi2"}; + Configurable useWeightedFinalPCA{"useWeightedFinalPCA", false, "Recalculate vertex position using track covariances, effective only if useAbsDCA is true"}; + Configurable maxR{"maxR", 200., "reject PCA's above this radius"}; + Configurable maxDZIni{"maxDZIni", 1e9, "reject (if>0) PCA candidate if tracks DZ exceeds threshold"}; + Configurable maxVtxChi2{"maxVtxChi2", 1e9, "reject (if>0) vtx. chi2 above this value"}; + Configurable minParamChange{"minParamChange", 1.e-3, "stop iterations if largest change of any X is smaller than this"}; + Configurable minRelChi2Change{"minRelChi2Change", 0.9, "stop iterations is chi2/chi2old > this"}; // Operation and minimisation criteria Configurable magneticField{"magneticField", 20.0f, "Magnetic field (in kilogauss)"}; Configurable doDCAplotsD{"doDCAplotsD", true, "do daughter prong DCA plots for D mesons"}; - Configurable doDCAplotsLc{"doDCAplotsLc", true, "do daughter prong DCA plots for Lc baryons"}; + Configurable doDCAplots3Prong{"doDCAplots3Prong", true, "do daughter prong DCA plots for Lc baryons"}; Configurable doTopoPlotsForSAndB{"doTopoPlotsForSAndB", true, "do topological variable distributions for S and B separately"}; Configurable mcSameMotherCheck{"mcSameMotherCheck", true, "check if tracks come from the same MC mother"}; Configurable dcaDaughtersSelection{"dcaDaughtersSelection", 1000.0f, "DCA between daughters (cm)"}; @@ -127,6 +140,13 @@ struct alice3decayFinder { o2::vertexing::DCAFitterN<2> fitter; o2::vertexing::DCAFitterN<3> fitter3; + double bz{0.}; + const float toMicrometers{10000.}; // from cm to µm + std::array daugsPdgCodes3Prong{{-1, -1, -1}}; + std::array daughtersMasses3Prong{{-1.f, -1.f, -1.f}}; + int motherPdgCode{-1}; + int charmHadFlag{0}; + HistogramRegistry histos{"histos", {}, OutputObjHandlingPolicy::AnalysisObject}; Partition trueD = aod::mcparticle::pdgCode == 421; @@ -149,30 +169,30 @@ struct alice3decayFinder { static constexpr uint32_t trackSelectionPrMinusFromLc = 1 << kInnerTOFProton | 1 << kOuterTOFProton | 1 << kRICHProton | 1 << kTruePrMinusFromLc; // partitions for D mesons - Partition tracksPiPlusFromD = + Partition tracksPiPlusFromD = ((aod::a3DecayMap::decayMap & trackSelectionPiPlusFromD) == trackSelectionPiPlusFromD) && aod::track::signed1Pt > 0.0f && nabs(aod::track::dcaXY) > piFromD_dcaXYconstant + piFromD_dcaXYpTdep* nabs(aod::track::signed1Pt); - Partition tracksPiMinusFromD = + Partition tracksPiMinusFromD = ((aod::a3DecayMap::decayMap & trackSelectionPiMinusFromD) == trackSelectionPiMinusFromD) && aod::track::signed1Pt < 0.0f && nabs(aod::track::dcaXY) > piFromD_dcaXYconstant + piFromD_dcaXYpTdep* nabs(aod::track::signed1Pt); - Partition tracksKaPlusFromD = + Partition tracksKaPlusFromD = ((aod::a3DecayMap::decayMap & trackSelectionKaPlusFromD) == trackSelectionKaPlusFromD) && aod::track::signed1Pt > 0.0f && nabs(aod::track::dcaXY) > kaFromD_dcaXYconstant + kaFromD_dcaXYpTdep* nabs(aod::track::signed1Pt); - Partition tracksKaMinusFromD = + Partition tracksKaMinusFromD = ((aod::a3DecayMap::decayMap & trackSelectionKaMinusFromD) == trackSelectionKaMinusFromD) && aod::track::signed1Pt < 0.0f && nabs(aod::track::dcaXY) > kaFromD_dcaXYconstant + kaFromD_dcaXYpTdep* nabs(aod::track::signed1Pt); // partitions for Lc baryons - Partition tracksPiPlusFromLc = + Partition tracksPiPlusFromLc = ((aod::a3DecayMap::decayMap & trackSelectionPiPlusFromLc) == trackSelectionPiPlusFromLc) && aod::track::signed1Pt > 0.0f && nabs(aod::track::dcaXY) > piFromLc_dcaXYconstant + piFromLc_dcaXYpTdep* nabs(aod::track::signed1Pt); - Partition tracksKaPlusFromLc = + Partition tracksKaPlusFromLc = ((aod::a3DecayMap::decayMap & trackSelectionKaPlusFromLc) == trackSelectionKaPlusFromLc) && aod::track::signed1Pt > 0.0f && nabs(aod::track::dcaXY) > kaFromLc_dcaXYconstant + kaFromLc_dcaXYpTdep* nabs(aod::track::signed1Pt); - Partition tracksPrPlusFromLc = + Partition tracksPrPlusFromLc = ((aod::a3DecayMap::decayMap & trackSelectionPrPlusFromLc) == trackSelectionPrPlusFromLc) && aod::track::signed1Pt > 0.0f && nabs(aod::track::dcaXY) > prFromLc_dcaXYconstant + prFromLc_dcaXYpTdep* nabs(aod::track::signed1Pt); // partitions for Lc baryons - Partition tracksPiMinusFromLc = + Partition tracksPiMinusFromLc = ((aod::a3DecayMap::decayMap & trackSelectionPiMinusFromLc) == trackSelectionPiMinusFromLc) && aod::track::signed1Pt < 0.0f && nabs(aod::track::dcaXY) > piFromLc_dcaXYconstant + piFromLc_dcaXYpTdep* nabs(aod::track::signed1Pt); - Partition tracksKaMinusFromLc = + Partition tracksKaMinusFromLc = ((aod::a3DecayMap::decayMap & trackSelectionKaMinusFromLc) == trackSelectionKaMinusFromLc) && aod::track::signed1Pt < 0.0f && nabs(aod::track::dcaXY) > kaFromLc_dcaXYconstant + kaFromLc_dcaXYpTdep* nabs(aod::track::signed1Pt); - Partition tracksPrMinusFromLc = + Partition tracksPrMinusFromLc = ((aod::a3DecayMap::decayMap & trackSelectionPrMinusFromLc) == trackSelectionPrMinusFromLc) && aod::track::signed1Pt < 0.0f && nabs(aod::track::dcaXY) > prFromLc_dcaXYconstant + prFromLc_dcaXYpTdep* nabs(aod::track::signed1Pt); // Helper struct to pass candidate information @@ -202,7 +222,30 @@ struct alice3decayFinder { float pt; float phi; float eta; - } lcbaryon; + std::array Pdaug0; // proton track + std::array Pdaug1; // kaon track + std::array Pdaug2; // pion track + std::array primaryVertex; // primary vertex coordinates + std::array secondaryVertex; // secondary vertex coordinates + float impactParameterY0; // impact parameters + float errorImpactParameterY0; // impact parameters error + float impactParameterY1; // impact parameters + float errorImpactParameterY1; // impact parameters error + float impactParameterY2; // impact parameters + float errorImpactParameterY2; // impact parameters error + float impactParameterZ0; // impact parameters + float errorImpactParameterZ0; // impact parameters error + float impactParameterZ1; // impact parameters + float errorImpactParameterZ1; // impact parameters error + float impactParameterZ2; // impact parameters + float errorImpactParameterZ2; // impact parameters error + float errorDecayLength; // normalized 3D decay length + float errorDecayLengthXY; // normalized 3D decay length + float chi2PCA; // normalized 3D decay length + int flagMc; // 0 = bkg, CharmHadAlice3 otherwise + int origin; // 1 = prompt, 2 = non-prompt + float ptBMotherRec; // pT of the B hadron mother (reconstructed) + } cand3prong; template bool buildDecayCandidateTwoBody(TTrackType const& posTrackRow, TTrackType const& negTrackRow, float posMass, float negMass, aod::McParticles const& mcParticles) @@ -272,18 +315,23 @@ struct alice3decayFinder { } template - bool buildDecayCandidateThreeBody(TTrackType const& prong0, TTrackType const& prong1, TTrackType const& prong2, float p0mass, float p1mass, float p2mass) + bool buildDecayCandidateThreeBody(aod::Collision const& collision, TTrackType const& prong0, TTrackType const& prong1, TTrackType const& prong2, aod::McParticles const& mcParticles) { - o2::track::TrackParCov t0 = getTrackParCov(prong0); - o2::track::TrackParCov t1 = getTrackParCov(prong1); - o2::track::TrackParCov t2 = getTrackParCov(prong2); + // get the collision primary vertex + auto primaryVertex = getPrimaryVertex(collision); + auto covMatrixPV = primaryVertex.getCov(); + + o2::track::TrackParCov trackParVar0 = getTrackParCov(prong0); + o2::track::TrackParCov trackParVar1 = getTrackParCov(prong1); + o2::track::TrackParCov trackParVar2 = getTrackParCov(prong2); //}-{}-{}-{}-{}-{}-{}-{}-{}-{} // Move close to minima int nCand = 0; try { - nCand = fitter3.process(t0, t1, t2); + nCand = fitter3.process(trackParVar0, trackParVar1, trackParVar2); } catch (...) { + LOG(info) << "Second vertex fit failed"; return false; } if (nCand == 0) { @@ -291,25 +339,102 @@ struct alice3decayFinder { } //}-{}-{}-{}-{}-{}-{}-{}-{}-{} - t0 = fitter3.getTrack(0); - t1 = fitter3.getTrack(1); - t2 = fitter3.getTrack(2); - std::array P0; - std::array P1; - std::array P2; - t0.getPxPyPzGlo(P0); - t1.getPxPyPzGlo(P1); - t2.getPxPyPzGlo(P2); - - lcbaryon.dcaDau = TMath::Sqrt(fitter3.getChi2AtPCACandidate()); - if (lcbaryon.dcaDau > dcaDaughtersSelection) + auto covMatrixPCA = fitter3.calcPCACovMatrixFlat(); + cand3prong.chi2PCA = fitter3.getChi2AtPCACandidate(); + cand3prong.dcaDau = TMath::Sqrt(fitter3.getChi2AtPCACandidate()); + if (cand3prong.dcaDau > dcaDaughtersSelection) { return false; + } + + cand3prong.primaryVertex = {primaryVertex.getX(), primaryVertex.getY(), primaryVertex.getZ()}; + auto secondaryVertex = fitter3.getPCACandidate(); + cand3prong.secondaryVertex = {secondaryVertex[0], secondaryVertex[1], secondaryVertex[2]}; + + trackParVar0 = fitter3.getTrack(0); + trackParVar1 = fitter3.getTrack(1); + trackParVar2 = fitter3.getTrack(2); + + std::array P0{}; + std::array P1{}; + std::array P2{}; + trackParVar0.getPxPyPzGlo(P0); + trackParVar1.getPxPyPzGlo(P1); + trackParVar2.getPxPyPzGlo(P2); + + o2::dataformats::DCA impactParameter0; + o2::dataformats::DCA impactParameter1; + o2::dataformats::DCA impactParameter2; + trackParVar0.propagateToDCA(primaryVertex, bz, &impactParameter0); + trackParVar1.propagateToDCA(primaryVertex, bz, &impactParameter1); + trackParVar2.propagateToDCA(primaryVertex, bz, &impactParameter2); + histos.fill(HIST("hDcaXYProngs"), prong0.pt(), impactParameter0.getY() * toMicrometers); + histos.fill(HIST("hDcaXYProngs"), prong1.pt(), impactParameter1.getY() * toMicrometers); + histos.fill(HIST("hDcaXYProngs"), prong2.pt(), impactParameter2.getY() * toMicrometers); + histos.fill(HIST("hDcaZProngs"), prong0.pt(), impactParameter0.getZ() * toMicrometers); + histos.fill(HIST("hDcaZProngs"), prong1.pt(), impactParameter1.getZ() * toMicrometers); + histos.fill(HIST("hDcaZProngs"), prong2.pt(), impactParameter2.getZ() * toMicrometers); + + // get uncertainty of the decay length + double phi, theta; + getPointDirection(std::array{primaryVertex.getX(), primaryVertex.getY(), primaryVertex.getZ()}, secondaryVertex, phi, theta); + cand3prong.errorDecayLength = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, theta) + getRotatedCovMatrixXX(covMatrixPCA, phi, theta)); + cand3prong.errorDecayLengthXY = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, 0.) + getRotatedCovMatrixXX(covMatrixPCA, phi, 0.)); + + cand3prong.impactParameterY0 = impactParameter0.getY(); + cand3prong.errorImpactParameterY0 = impactParameter0.getSigmaY2(); + cand3prong.impactParameterY1 = impactParameter1.getY(); + cand3prong.errorImpactParameterY1 = impactParameter1.getSigmaY2(); + cand3prong.impactParameterY2 = impactParameter2.getY(); + cand3prong.errorImpactParameterY2 = impactParameter2.getSigmaY2(); + + cand3prong.impactParameterZ0 = impactParameter0.getZ(); + cand3prong.errorImpactParameterZ0 = impactParameter0.getSigmaZ2(); + cand3prong.impactParameterZ1 = impactParameter1.getZ(); + cand3prong.errorImpactParameterZ1 = impactParameter1.getSigmaZ2(); + cand3prong.impactParameterZ2 = impactParameter2.getZ(); + cand3prong.errorImpactParameterZ2 = impactParameter2.getSigmaZ2(); // return mass - lcbaryon.mass = RecoDecay::m(array{array{P0[0], P0[1], P0[2]}, array{P1[0], P1[1], P1[2]}, array{P2[0], P2[1], P2[2]}}, array{p0mass, p1mass, p2mass}); - lcbaryon.pt = std::hypot(P0[0] + P1[0] + P2[0], P0[1] + P1[1] + P2[1]); - lcbaryon.phi = RecoDecay::phi(array{P0[0] + P1[0] + P2[0], P0[1] + P1[1] + P2[1]}); - lcbaryon.eta = RecoDecay::eta(array{P0[0] + P1[0] + P2[0], P0[1] + P1[1] + P2[1], P0[2] + P1[2] + P2[2]}); + cand3prong.mass = RecoDecay::m(array{array{P0[0], P0[1], P0[2]}, + array{P1[0], P1[1], P1[2]}, + array{P2[0], P2[1], P2[2]}}, + daughtersMasses3Prong); + + cand3prong.pt = std::hypot(P0[0] + P1[0] + P2[0], P0[1] + P1[1] + P2[1]); + cand3prong.phi = RecoDecay::phi(array{P0[0] + P1[0] + P2[0], P0[1] + P1[1] + P2[1]}); + cand3prong.eta = RecoDecay::eta(array{P0[0] + P1[0] + P2[0], P0[1] + P1[1] + P2[1], P0[2] + P1[2] + P2[2]}); + cand3prong.Pdaug0[0] = P0[0]; + cand3prong.Pdaug0[1] = P0[1]; + cand3prong.Pdaug0[2] = P0[2]; + cand3prong.Pdaug1[0] = P1[0]; + cand3prong.Pdaug1[1] = P1[1]; + cand3prong.Pdaug1[2] = P1[2]; + cand3prong.Pdaug2[0] = P2[0]; + cand3prong.Pdaug2[1] = P2[1]; + cand3prong.Pdaug2[2] = P2[2]; + + // MC truth check + cand3prong.flagMc = 0; // bkg + int8_t sign = 0; + auto arrayDaughters = std::array{prong0, prong1, prong2}; + int indexRec = RecoDecay::getMatchedMCRec(mcParticles, arrayDaughters, motherPdgCode, daugsPdgCodes3Prong, true, &sign, 2); + auto motherPart = mcParticles.rawIteratorAt(indexRec); + if (indexRec > -1) { + cand3prong.flagMc = motherPart.pdgCode() > 0 ? charmHadFlag : -charmHadFlag; // Particle + } + + cand3prong.origin = 0; + if (indexRec > -1) { + auto motherParticle = mcParticles.rawIteratorAt(indexRec); + std::vector idxBhadMothers{}; + int origin = RecoDecay::getCharmHadronOrigin(mcParticles, motherParticle, false, &idxBhadMothers); + cand3prong.origin = origin; + cand3prong.ptBMotherRec = -1.f; + if (origin == RecoDecay::OriginType::NonPrompt) { + auto bHadMother = mcParticles.rawIteratorAt(idxBhadMothers[0]); + cand3prong.ptBMotherRec = bHadMother.pt(); + } + } return true; } @@ -339,25 +464,25 @@ struct alice3decayFinder { void init(InitContext&) { // initialize O2 2-prong fitter (only once) - fitter.setPropagateToPCA(true); - fitter.setMaxR(200.); - fitter.setMinParamChange(1e-3); - fitter.setMinRelChi2Change(0.9); - fitter.setMaxDZIni(1e9); - fitter.setMaxChi2(1e9); - fitter.setUseAbsDCA(true); - fitter.setWeightedFinalPCA(false); + fitter.setPropagateToPCA(propagateToPCA); + fitter.setMaxR(maxR); + fitter.setMinParamChange(minParamChange); + fitter.setMinRelChi2Change(minRelChi2Change); + fitter.setMaxDZIni(maxDZIni); + fitter.setMaxChi2(maxVtxChi2); + fitter.setUseAbsDCA(useAbsDCA); + fitter.setWeightedFinalPCA(useWeightedFinalPCA); fitter.setBz(magneticField); fitter.setMatCorrType(o2::base::Propagator::MatCorrType::USEMatCorrNONE); - fitter3.setPropagateToPCA(true); - fitter3.setMaxR(200.); - fitter3.setMinParamChange(1e-3); - fitter3.setMinRelChi2Change(0.9); - fitter3.setMaxDZIni(1e9); - fitter3.setMaxChi2(1e9); - fitter3.setUseAbsDCA(true); - fitter3.setWeightedFinalPCA(false); + fitter3.setPropagateToPCA(propagateToPCA); + fitter3.setMaxR(maxR); + fitter3.setMinParamChange(minParamChange); + fitter3.setMinRelChi2Change(minRelChi2Change); + fitter3.setMaxDZIni(maxDZIni); + fitter3.setMaxChi2(maxVtxChi2); + fitter3.setUseAbsDCA(useAbsDCA); + fitter3.setWeightedFinalPCA(useWeightedFinalPCA); fitter3.setBz(magneticField); fitter3.setMatCorrType(o2::base::Propagator::MatCorrType::USEMatCorrNONE); @@ -445,30 +570,37 @@ struct alice3decayFinder { } } } - if (doprocessFindLcBaryons) { - histos.add("h2dGenLc", "h2dGenLc", kTH2F, {axisPt, axisEta}); - histos.add("h2dGenLcbar", "h2dGenLcbar", kTH2F, {axisPt, axisEta}); - histos.add("h3dRecLc", "h2dRecLc", kTH3F, {axisPt, axisEta, axisLcMass}); - histos.add("h3dRecLcbar", "h2dRecLcbar", kTH3F, {axisPt, axisEta, axisLcMass}); - - histos.add("hMassLc", "hMassLc", kTH1F, {axisLcMass}); - histos.add("hMassLcbar", "hMassLcbar", kTH1F, {axisLcMass}); - - if (doDCAplotsD) { - histos.add("hDCALcDaughters", "hDCALcDaughters", kTH1D, {axisDCADaughters}); - histos.add("hDCALcbarDaughters", "hDCALcbarDaughters", kTH1D, {axisDCA}); - histos.add("h2dDCAxyVsPtPiPlusFromLc", "h2dDCAxyVsPtPiPlusFromLc", kTH2F, {axisPt, axisDCA}); - histos.add("h2dDCAxyVsPtPiMinusFromLc", "h2dDCAxyVsPtPiMinusFromLc", kTH2F, {axisPt, axisDCA}); - histos.add("h2dDCAxyVsPtKaPlusFromLc", "h2dDCAxyVsPtKaPlusFromLc", kTH2F, {axisPt, axisDCA}); - histos.add("h2dDCAxyVsPtKaMinusFromLc", "h2dDCAxyVsPtKaMinusFromLc", kTH2F, {axisPt, axisDCA}); - histos.add("h2dDCAxyVsPtPrPlusFromLc", "h2dDCAxyVsPtPrPlusFromLc", kTH2F, {axisPt, axisDCA}); - histos.add("h2dDCAxyVsPtPrMinusFromLc", "h2dDCAxyVsPtPrMinusFromLc", kTH2F, {axisPt, axisDCA}); + if (doprocessFindLc) { + histos.add("h2dGen3Prong", "h2dGen3Prong", kTH2F, {axisPt, axisEta}); + histos.add("h2dGen3ProngBar", "h2dGen3ProngBar", kTH2F, {axisPt, axisEta}); + histos.add("h3dRec3Prong", "h3dRec3Prong", kTH3F, {axisPt, axisEta, axisLcMass}); + histos.add("hMass3Prong", "hMass3Prong", kTH1F, {axisLcMass}); + + if (doDCAplots3Prong) { + histos.add("hDCA3ProngDaughters", "hDCA3ProngDaughters", kTH1D, {axisDCADaughters}); + histos.add("h2dDCAxyVsPtPiPlusFrom3P", "h2dDCAxyVsPtPiPlusFrom3P", kTH2F, {axisPt, axisDCA}); + histos.add("h2dDCAxyVsPtPiMinusFrom3P", "h2dDCAxyVsPtPiMinusFrom3P", kTH2F, {axisPt, axisDCA}); + histos.add("h2dDCAxyVsPtKaPlusFrom3P", "h2dDCAxyVsPtKaPlusFrom3P", kTH2F, {axisPt, axisDCA}); + histos.add("h2dDCAxyVsPtKaMinusFrom3P", "h2dDCAxyVsPtKaMinusFrom3P", kTH2F, {axisPt, axisDCA}); + histos.add("h2dDCAxyVsPtPrPlusFrom3P", "h2dDCAxyVsPtPrPlusFrom3P", kTH2F, {axisPt, axisDCA}); + histos.add("h2dDCAxyVsPtPrMinusFrom3P", "h2dDCAxyVsPtPrMinusFrom3P", kTH2F, {axisPt, axisDCA}); + histos.add("hDcaXYProngs", "DCAxy of 3-prong candidate daughters;#it{p}_{T} (GeV/#it{c};#it{d}_{xy}) (#mum);entries", {HistType::kTH2F, {{100, 0., 20.}, {200, -500., 500.}}}); + histos.add("hDcaZProngs", "DCAz of 3-prong candidate daughters;#it{p}_{T} (GeV/#it{c};#it{d}_{z}) (#mum);entries", {HistType::kTH2F, {{100, 0., 20.}, {200, -500., 500.}}}); } } + + if (doprocessFindLc) { + daugsPdgCodes3Prong = {+kProton, -kKPlus, +kPiPlus}; + motherPdgCode = o2::constants::physics::Pdg::kLambdaCPlus; + daughtersMasses3Prong = {o2::constants::physics::MassProton, + o2::constants::physics::MassKaonCharged, + o2::constants::physics::MassPionCharged}; + charmHadFlag = CharmHadAlice3::Lc; + } } //*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+* - void processGenerated(aod::McParticles const&) + void processGenerated(aod::McParticles const& mcParticles) { // no grouping for MC particles -> as intended if (doprocessFindDmesons) { @@ -507,16 +639,31 @@ struct alice3decayFinder { } } } - if (doprocessFindLcBaryons) { - for (auto const& mcParticle : trueLc) - histos.fill(HIST("h2dGenLc"), mcParticle.pt(), mcParticle.eta()); - for (auto const& mcParticle : trueLcbar) - histos.fill(HIST("h2dGenLcbar"), mcParticle.pt(), mcParticle.eta()); + if (doprocessFindLc) { + for (auto const& mcParticle : mcParticles) { + if (std::abs(mcParticle.pdgCode()) != motherPdgCode) { + mcGenFlags(-1, -1, -1); + continue; + } + std::vector idxBhadMothers{}; + int origin = RecoDecay::getCharmHadronOrigin(mcParticles, mcParticle, false, &idxBhadMothers); + float ptBMotherGen{-1.f}; + if (origin == RecoDecay::OriginType::NonPrompt) { + auto bHadMother = mcParticles.rawIteratorAt(idxBhadMothers[0]); + ptBMotherGen = bHadMother.pt(); + } + mcGenFlags(origin, ptBMotherGen, mcParticle.pdgCode() ? charmHadFlag : -charmHadFlag); + if (mcParticle.pdgCode() > 0) { + histos.fill(HIST("h2dGen3Prong"), mcParticle.pt(), mcParticle.eta()); + } else { + histos.fill(HIST("h2dGen3ProngBar"), mcParticle.pt(), mcParticle.eta()); + } + } } } //*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+* - void processFindDmesons(aod::Collision const& collision, alice3tracks const&, aod::McParticles const& mcParticles) + void processFindDmesons(aod::Collision const& collision, Alice3TracksWPid const&, aod::McParticles const& mcParticles) { // group with this collision auto tracksPiPlusFromDgrouped = tracksPiPlusFromD->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); @@ -635,7 +782,7 @@ struct alice3decayFinder { continue; if (decayLengthXY < DMinDecayLengthXY || decayLengthXY > DMaxDecayLengthXY) continue; - auto decayLengthSquaredCut = std::min((std::hypot(dmeson.P[0], dmeson.P[1], dmeson.P[2]) * 0.0066) + 0.01, (double)DDecayLengthSquaredCut); + auto decayLengthSquaredCut = std::min((std::hypot(dmeson.P[0], dmeson.P[1], dmeson.P[2]) * 0.0066) + 0.01, static_cast(DDecayLengthSquaredCut)); if (decayLength * decayLength < decayLengthSquaredCut * decayLengthSquaredCut) continue; @@ -782,7 +929,7 @@ struct alice3decayFinder { continue; if (decayLengthXY < DMinDecayLengthXY || decayLengthXY > DMaxDecayLengthXY) continue; - auto decayLengthSquaredCut = std::min((std::hypot(dmeson.P[0], dmeson.P[1], dmeson.P[2]) * 0.0066) + 0.01, (double)DDecayLengthSquaredCut); + auto decayLengthSquaredCut = std::min((std::hypot(dmeson.P[0], dmeson.P[1], dmeson.P[2]) * 0.0066) + 0.01, static_cast(DDecayLengthSquaredCut)); if (decayLength * decayLength < decayLengthSquaredCut * decayLengthSquaredCut) continue; @@ -831,62 +978,63 @@ struct alice3decayFinder { } //*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+* - //*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+* - void processFindLcBaryons(aod::Collision const& collision, alice3tracks const&, aod::McParticles const&) + template + void fillPidTable(TProng const& prong0, TProng const& prong1, TProng const& prong2) { - // group with this collision - auto tracksPiPlusFromLcgrouped = tracksPiPlusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); - auto tracksKaPlusFromLcgrouped = tracksKaPlusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); - auto tracksPrPlusFromLcgrouped = tracksPrPlusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); - - auto tracksPiMinusFromLcgrouped = tracksPiMinusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); - auto tracksKaMinusFromLcgrouped = tracksKaMinusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); - auto tracksPrMinusFromLcgrouped = tracksPrMinusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); - - if (doDCAplotsLc) { - for (auto const& track : tracksPiPlusFromLcgrouped) - histos.fill(HIST("h2dDCAxyVsPtPiPlusFromLc"), track.pt(), track.dcaXY() * 1e+4); - for (auto const& track : tracksPiMinusFromLcgrouped) - histos.fill(HIST("h2dDCAxyVsPtPiMinusFromLc"), track.pt(), track.dcaXY() * 1e+4); - for (auto const& track : tracksKaPlusFromLcgrouped) - histos.fill(HIST("h2dDCAxyVsPtKaPlusFromLc"), track.pt(), track.dcaXY() * 1e+4); - for (auto const& track : tracksKaMinusFromLcgrouped) - histos.fill(HIST("h2dDCAxyVsPtKaMinusFromLc"), track.pt(), track.dcaXY() * 1e+4); - for (auto const& track : tracksPrPlusFromLcgrouped) - histos.fill(HIST("h2dDCAxyVsPtPrPlusFromLc"), track.pt(), track.dcaXY() * 1e+4); - for (auto const& track : tracksPrMinusFromLcgrouped) - histos.fill(HIST("h2dDCAxyVsPtPrMinusFromLc"), track.pt(), track.dcaXY() * 1e+4); + if (motherPdgCode == o2::constants::physics::Pdg::kLambdaCPlus) { + pidInfoLcDaugs(prong0.nSigmaTrkPr(), prong0.nSigmaProtonRich(), prong0.nSigmaProtonInnerTOF(), prong0.nSigmaProtonOuterTOF(), + prong1.nSigmaTrkKa(), prong1.nSigmaKaonRich(), prong1.nSigmaKaonInnerTOF(), prong1.nSigmaKaonOuterTOF(), + prong2.nSigmaTrkPi(), prong2.nSigmaPionRich(), prong2.nSigmaPionInnerTOF(), prong2.nSigmaPionOuterTOF()); + } else { + LOG(fatal) << "3-prong candidate not implemented yet"; } + } - // Lc+ baryons +4122 -> +2212 -321 +211 - for (auto const& proton : tracksPrPlusFromLcgrouped) { - for (auto const& pion : tracksPiPlusFromLcgrouped) { - if (pion.globalIndex() == proton.globalIndex()) - continue; // avoid self - for (auto const& kaon : tracksKaMinusFromLcgrouped) { - if (mcSameMotherCheck && (!checkSameMother(proton, kaon) || !checkSameMother(proton, pion))) - continue; - if (!buildDecayCandidateThreeBody(proton, kaon, pion, o2::constants::physics::MassProton, o2::constants::physics::MassKaonCharged, o2::constants::physics::MassPionCharged)) - continue; - histos.fill(HIST("hDCALcDaughters"), lcbaryon.dcaDau * 1e+4); - histos.fill(HIST("hMassLc"), lcbaryon.mass); - histos.fill(HIST("h3dRecLc"), lcbaryon.pt, lcbaryon.eta, lcbaryon.mass); - } - } - } - // Lc- baryons -4122 -> -2212 +321 -211 - for (auto const& proton : tracksPrMinusFromLcgrouped) { - for (auto const& pion : tracksPiMinusFromLcgrouped) { - if (pion.globalIndex() == proton.globalIndex()) + //*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+* + template + void fill3ProngTable(aod::Collision const& collision, TProng const& prongs0, TProng const& prongs1, TProng const& prongs2, aod::McParticles const& mcParticles) + { + for (auto const& prong0 : prongs0) { + for (auto const& prong2 : prongs2) { + if (prong2.globalIndex() == prong0.globalIndex()) continue; // avoid self - for (auto const& kaon : tracksKaPlusFromLcgrouped) { - if (mcSameMotherCheck && (!checkSameMother(proton, kaon) || !checkSameMother(proton, pion))) + for (auto const& prong1 : prongs1) { + if (mcSameMotherCheck && (!checkSameMother(prong0, prong1) || !checkSameMother(prong0, prong1))) { continue; - if (!buildDecayCandidateThreeBody(proton, kaon, pion, o2::constants::physics::MassProton, o2::constants::physics::MassKaonCharged, o2::constants::physics::MassPionCharged)) + } + if (!buildDecayCandidateThreeBody(collision, prong0, prong1, prong2, mcParticles)) { continue; - histos.fill(HIST("hDCALcbarDaughters"), lcbaryon.dcaDau * 1e+4); - histos.fill(HIST("hMassLcbar"), lcbaryon.mass); - histos.fill(HIST("h3dRecLcbar"), lcbaryon.pt, lcbaryon.eta, lcbaryon.mass); + } + histos.fill(HIST("hDCA3ProngDaughters"), cand3prong.dcaDau * 1e+4); + histos.fill(HIST("hMass3Prong"), cand3prong.mass); + histos.fill(HIST("h3dRec3Prong"), cand3prong.pt, cand3prong.eta, cand3prong.mass); + + auto candPx = cand3prong.Pdaug0[0] + cand3prong.Pdaug1[0] + cand3prong.Pdaug2[0]; + auto candPy = cand3prong.Pdaug0[1] + cand3prong.Pdaug1[1] + cand3prong.Pdaug2[1]; + auto candPz = cand3prong.Pdaug0[2] + cand3prong.Pdaug1[2] + cand3prong.Pdaug2[2]; + + candidate3Prong(collision.globalIndex(), + cand3prong.primaryVertex[0], cand3prong.primaryVertex[1], cand3prong.primaryVertex[2], + cand3prong.secondaryVertex[0], cand3prong.secondaryVertex[1], cand3prong.secondaryVertex[2], + cand3prong.errorDecayLength, cand3prong.errorDecayLengthXY, + cand3prong.chi2PCA, + cand3prong.eta, + cand3prong.phi, + cand3prong.pt, + cand3prong.Pdaug2[0], cand3prong.Pdaug2[1], cand3prong.Pdaug2[2], + cand3prong.Pdaug1[0], cand3prong.Pdaug1[1], cand3prong.Pdaug1[2], + cand3prong.Pdaug0[0], cand3prong.Pdaug0[1], cand3prong.Pdaug0[2], + cand3prong.impactParameterY0, cand3prong.impactParameterY1, cand3prong.impactParameterY2, + std::sqrt(cand3prong.errorImpactParameterY0), + std::sqrt(cand3prong.errorImpactParameterY1), + std::sqrt(cand3prong.errorImpactParameterY2), + cand3prong.impactParameterZ0, cand3prong.impactParameterZ1, cand3prong.impactParameterZ2, + std::sqrt(cand3prong.errorImpactParameterZ0), + std::sqrt(cand3prong.errorImpactParameterZ1), + std::sqrt(cand3prong.errorImpactParameterZ2), + candPx, candPy, candPz); + mcRecFlags(cand3prong.origin, cand3prong.ptBMotherRec, cand3prong.flagMc); // placeholder for prompt/non-prompt + fillPidTable(prong0, prong1, prong2); } } } @@ -896,12 +1044,44 @@ struct alice3decayFinder { //*>-~-<*>-~-<*>-~-<*>-~-<*>-~-<*>-~-<*>-~-<*>-~-<* PROCESS_SWITCH(alice3decayFinder, processGenerated, "fill MC-only histograms", true); PROCESS_SWITCH(alice3decayFinder, processFindDmesons, "find D mesons", true); - PROCESS_SWITCH(alice3decayFinder, processFindLcBaryons, "find Lc Baryons", true); + + void processFindLc(aod::Collision const& collision, + aod::McParticles const& mcParticles, + Alice3TracksWPid const&) + { + + auto tracksPiPlus = tracksPiPlusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); + auto tracksKaPlus = tracksKaPlusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); + auto tracksPrPlus = tracksPrPlusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); + auto tracksPiMinus = tracksPiMinusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); + auto tracksKaMinus = tracksKaMinusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); + auto tracksPrMinus = tracksPrMinusFromLc->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache); + + if (doDCAplots3Prong) { + for (auto const& track : tracksPiPlus) + histos.fill(HIST("h2dDCAxyVsPtPiPlusFrom3P"), track.pt(), track.dcaXY() * 1e+4); + for (auto const& track : tracksPiMinus) + histos.fill(HIST("h2dDCAxyVsPtPiMinusFrom3P"), track.pt(), track.dcaXY() * 1e+4); + for (auto const& track : tracksKaPlus) + histos.fill(HIST("h2dDCAxyVsPtKaPlusFrom3P"), track.pt(), track.dcaXY() * 1e+4); + for (auto const& track : tracksKaMinus) + histos.fill(HIST("h2dDCAxyVsPtKaMinusFrom3P"), track.pt(), track.dcaXY() * 1e+4); + for (auto const& track : tracksPrPlus) + histos.fill(HIST("h2dDCAxyVsPtPrPlusFrom3P"), track.pt(), track.dcaXY() * 1e+4); + for (auto const& track : tracksPrMinus) + histos.fill(HIST("h2dDCAxyVsPtPrMinusFrom3P"), track.pt(), track.dcaXY() * 1e+4); + } + + // Particles + fill3ProngTable(collision, tracksPrPlus, tracksKaMinus, tracksPiPlus, mcParticles); + fill3ProngTable(collision, tracksPrMinus, tracksKaPlus, tracksPiMinus, mcParticles); + } + PROCESS_SWITCH(alice3decayFinder, processFindLc, "find Lc Baryons", true); //*>-~-<*>-~-<*>-~-<*>-~-<*>-~-<*>-~-<*>-~-<*>-~-<* }; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) { return WorkflowSpec{ - adaptAnalysisTask(cfgc)}; + adaptAnalysisTask(cfgc)}; // o2-linter: disable=name/o2-task (wrong hyphenation) } diff --git a/ALICE3/TableProducer/alice3HfSelector3Prong.cxx b/ALICE3/TableProducer/alice3HfSelector3Prong.cxx new file mode 100644 index 00000000000..4fe0f4d1e19 --- /dev/null +++ b/ALICE3/TableProducer/alice3HfSelector3Prong.cxx @@ -0,0 +1,380 @@ +// 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 alice3HfSelector3prong.cxx +/// \brief 3-prong candidates selection task +/// +/// \author Marcello Di Costanzo , Polytechnic University of Turin and INFN Turin + +#include "PWGHF/Core/SelectorCuts.h" +#include "PWGHF/DataModel/CandidateSelectionTables.h" + +#include "ALICE3/DataModel/A3DecayFinderTables.h" +#include "ALICE3/DataModel/OTFPIDTrk.h" +#include "ALICE3/DataModel/OTFRICH.h" +#include "ALICE3/DataModel/OTFTOF.h" +#include "ALICE3/DataModel/RICH.h" +#include "ALICE3/ML/HfMlResponse3Prong.h" +#include "ALICE3/Utils/utilsHfAlice3.h" +#include "ALICE3/Utils/utilsSelectionsAlice3.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +using namespace o2; +using namespace o2::analysis; +using namespace o2::framework; +using namespace o2::aod::a3_hf_sel_3prong; + +/// Struct for applying Lc selection cuts +struct Alice3HfSelector3Prong { + Produces candSelFlags; // flags for isSelLc + Produces candMlScores; + + Configurable ptCandMin{"ptCandMin", 0., "Lower bound of cand pT"}; + Configurable ptCandMax{"ptCandMax", 36., "Upper bound of cand pT"}; + // TRK PID + Configurable ptPidTrkMin{"ptPidTrkMin", 0.1, "Lower bound of track pT for Trk PID"}; + Configurable ptPidTrkMax{"ptPidTrkMax", 1., "Upper bound of track pT for Trk PID"}; + Configurable nSigmaTrkMax{"nSigmaTrkMax", 3., "Nsigma cut on Trk only"}; + // RICH PID + Configurable ptPidRichMin{"ptPidRichMin", 0.1, "Lower bound of track pT for Rich PID"}; + Configurable ptPidRichMax{"ptPidRichMax", 1., "Upper bound of track pT for Rich PID"}; + Configurable nSigmaRichMax{"nSigmaRichMax", 3., "Nsigma cut on Rich only"}; + // INNER TOF PID + Configurable ptPidInnTofMin{"ptPidInnTofMin", 0.5, "Lower bound of track pT for InTOF PID"}; + Configurable ptPidInnTofMax{"ptPidInnTofMax", 2.5, "Upper bound of track pT for InTOF PID"}; + Configurable nSigmaInnTofMax{"nSigmaInnTofMax", 3., "Nsigma cut on InTOF only"}; + // OUTER TOF PID + Configurable ptPidOutTofMin{"ptPidOutTofMin", 0.5, "Lower bound of track pT for OutTOF PID"}; + Configurable ptPidOutTofMax{"ptPidOutTofMax", 2.5, "Upper bound of track pT for OutTOF PID"}; + Configurable nSigmaOutTofMax{"nSigmaOutTofMax", 3., "Nsigma cut on OutTOF only"}; + // DCA track cuts + Configurable> binsPtTrack{"binsPtTrack", std::vector{hf_cuts_single_track::vecBinsPtTrack}, "track pT bin limits for DCA XY/Z pT-dependent cut"}; + Configurable> cutsSingleTrack{"cutsSingleTrack", {hf_cuts_single_track::CutsTrack[0], hf_cuts_single_track::NBinsPtTrack, hf_cuts_single_track::NCutVarsTrack, hf_cuts_single_track::labelsPtTrack, hf_cuts_single_track::labelsCutVarTrack}, "Single-track selections"}; + // topological cuts + Configurable> binsPt{"binsPt", std::vector{hf_cuts_3prongs_alice3::vecBinsPt}, "pT bin limits"}; + Configurable> cuts{"cuts", {hf_cuts_3prongs_alice3::Cuts[0], hf_cuts_3prongs_alice3::NBinsPt, hf_cuts_3prongs_alice3::NCutVars, hf_cuts_3prongs_alice3::labelsPt, hf_cuts_3prongs_alice3::labelsCutVar}, "Lc cand selection per pT bin"}; + // QA switch + Configurable activateQA{"activateQA", false, "Flag to enable QA histogram"}; + // ML inference + Configurable applyMl{"applyMl", false, "Flag to apply ML selections"}; + Configurable> binsPtMl{"binsPtMl", std::vector{hf_cuts_ml::vecBinsPt}, "pT bin limits for ML application"}; + Configurable> cutDirMl{"cutDirMl", std::vector{hf_cuts_ml::vecCutDir}, "Whether to reject score values greater or smaller than the threshold"}; + Configurable> cutsMl{"cutsMl", {hf_cuts_ml::Cuts[0], hf_cuts_ml::NBinsPt, hf_cuts_ml::NCutScores, hf_cuts_ml::labelsPt, hf_cuts_ml::labelsCutScore}, "ML selections per pT bin"}; + Configurable nClassesMl{"nClassesMl", static_cast(hf_cuts_ml::NCutScores), "Number of classes in ML model"}; + Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"feature1", "feature2"}, "Names of ML model input features"}; + // CCDB configuration + Configurable ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; + Configurable> modelPathsCCDB{"modelPathsCCDB", std::vector{"EventFiltering/PWGHF/BDTLc"}, "Paths of models on CCDB"}; + Configurable> onnxFileNames{"onnxFileNames", std::vector{"ModelHandler_onnx_MassHypo0.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"}; + Configurable timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB"}; + Configurable loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; + + HfHelperAlice3 hfHelper; + o2::analysis::HfMlResponse3Prong mlResponse; + o2::ccdb::CcdbApi ccdbApi; + + using CandsLc = soa::Join; + using CandsLcWMcTruth = soa::Join; + + float MassReference{-1.f}; + + HistogramRegistry registry{"registry"}; + + void init(InitContext const&) + { + if (activateQA) { + constexpr int kNBinsSelections = 1 + aod::SelectionStep::NSelectionSteps; + std::string labels[kNBinsSelections]; + labels[0] = "No selection"; + labels[1 + aod::SelectionStep::RecoSkims] = "Skims selection"; + labels[1 + aod::SelectionStep::RecoTopol] = "Skims & Topological selections"; + labels[1 + aod::SelectionStep::RecoPID] = "Skims & Topological & PID selections"; + labels[1 + aod::SelectionStep::RecoMl] = "ML selection"; + static const AxisSpec axisSelections = {kNBinsSelections, 0.5, kNBinsSelections + 0.5, ""}; + registry.add("hSelections", "Selections;;#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {axisSelections, {(std::vector)binsPt, "#it{p}_{T} (GeV/#it{c})"}}}); + for (int iBin = 0; iBin < kNBinsSelections; ++iBin) { + registry.get(HIST("hSelections"))->GetXaxis()->SetBinLabel(iBin + 1, labels[iBin].data()); + } + } + + if (applyMl) { + mlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); + if (loadModelsFromCCDB) { + ccdbApi.init(ccdbUrl); + mlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB); + } else { + mlResponse.setModelPathsLocal(onnxFileNames); + } + mlResponse.cacheInputFeaturesIndices(namesInputFeatures); + mlResponse.init(); + } + + if (doprocessLc) { + MassReference = o2::constants::physics::MassLambdaCPlus; + } + } + + /// Conjugate-independent topological cuts + /// \tparam T is candidate type + /// \param cand is cand + /// \param candPt is candidate pT + /// \return true if cand passes all cuts + template + bool selectionTopol(const T& cand, float candPt) + { + int const ptBin = findBin(binsPt, candPt); + + // check that the cand pT is within the analysis range + if (candPt < ptCandMin || candPt >= ptCandMax) { + return false; + } + + // cut on daughter pT + if (cand.ptProng0() < cuts->get(ptBin, "pT prong 0") || + cand.ptProng1() < cuts->get(ptBin, "pT prong 1") || + cand.ptProng2() < cuts->get(ptBin, "pT prong 2")) { + return false; + } + + // cosine of pointing angle + if (cand.cpa() <= cuts->get(ptBin, "cos pointing angle")) { + return false; + } + + // cand chi2PCA + if (cand.chi2PCA() > cuts->get(ptBin, "Chi2PCA")) { + return false; + } + + if (cand.decayLength() <= cuts->get(ptBin, "decay length")) { + return false; + } + + // cand decay length XY + if (cand.decayLengthXY() <= cuts->get(ptBin, "decLengthXY")) { + return false; + } + + // cand normalized decay length XY + if (cand.decayLengthXYNormalised() < cuts->get(ptBin, "normDecLXY")) { + return false; + } + + // cand impact parameter XY + if (std::abs(cand.impactParameterXY()) > cuts->get(ptBin, "impParXY")) { + return false; + } + + // cand daughter prong DCA + if (!isSelectedCandidateProngDca(cand)) { + return false; + } + + return true; + } + + /// Candidate mass selection + /// \tparam CharmHad is the charm hadron type + /// \tparam SwapHypo indicates whether to swap mass hypothesis or not + /// \tparam TCandidate is candidate type + /// \param ptBin is candidate pT bin + /// \param cand is candidate + /// \return true if candidate passes mass selection + template + bool selectionCandidateMass(int const ptBin, const TCandidate& cand) + { + float massCand = hfHelper.getCandMass(cand); + // cut on mass window + if (std::abs(massCand - MassReference) > cuts->get(ptBin, "m")) { + return false; + } + + return true; + } + + /// Single-track dca_xy and dca_z cuts + /// \tparam T1 is candidate type + /// \param cand is the Lc cand + /// \return true if all the prongs pass the selections + template + bool isSelectedCandidateProngDca(const T1& cand) + { + return (isSelectedTrackDca(binsPtTrack, cutsSingleTrack, cand.ptProng0(), cand.impactParameterY0(), cand.impactParameterZ0()) && + isSelectedTrackDca(binsPtTrack, cutsSingleTrack, cand.ptProng1(), cand.impactParameterY1(), cand.impactParameterZ1()) && + isSelectedTrackDca(binsPtTrack, cutsSingleTrack, cand.ptProng2(), cand.impactParameterY2(), cand.impactParameterZ2())); + } + + /// Apply PID selection + /// \tparam CharmHad is charm hadron type + /// \tparam TCand is candidate type + /// \param cand is candidate + /// \param pidMask is bitmask to be configured + template + void configurePidMask(const TCand& cand, uint32_t& pidMask) + { + + auto isSelPid = [&](int selCut, float nsigma, float pt, float nSigmaMax, float ptMin, float ptMax) { + bool isSelected = !(pt >= ptMin && pt < ptMax && std::abs(nsigma) > nSigmaMax); + if (isSelected) + SETBIT(pidMask, selCut); + return isSelected; + }; + + // prong 0 + float ptProng0{cand.ptProng0()}; + if constexpr (CharmHad == CharmHadAlice3::Lc) { + isSelPid(cand.nSigTrkPr0(), PidSels::TrkProng0, ptProng0, nSigmaTrkMax, ptPidTrkMin, ptPidTrkMax); + isSelPid(cand.nSigRichPr0(), PidSels::RichProng0, ptProng0, nSigmaRichMax, ptPidRichMin, ptPidRichMax); + isSelPid(cand.nSigInnTofPr0(), PidSels::InnTofProng0, ptProng0, nSigmaInnTofMax, ptPidInnTofMin, ptPidInnTofMax); + isSelPid(cand.nSigOutTofPr0(), PidSels::OutTofProng0, ptProng0, nSigmaOutTofMax, ptPidOutTofMin, ptPidOutTofMax); + } + + // prong 1 + float ptProng1{cand.ptProng1()}; + if constexpr (CharmHad == CharmHadAlice3::Lc) { + isSelPid(cand.nSigTrkKa1(), PidSels::TrkProng1, ptProng1, nSigmaTrkMax, ptPidTrkMin, ptPidTrkMax); + isSelPid(cand.nSigRichKa1(), PidSels::RichProng1, ptProng1, nSigmaRichMax, ptPidRichMin, ptPidRichMax); + isSelPid(cand.nSigInnTofKa1(), PidSels::InnTofProng1, ptProng1, nSigmaInnTofMax, ptPidInnTofMin, ptPidInnTofMax); + isSelPid(cand.nSigOutTofKa1(), PidSels::OutTofProng1, ptProng1, nSigmaOutTofMax, ptPidOutTofMin, ptPidOutTofMax); + } + + // prong 2 + float ptProng2{cand.ptProng2()}; + if constexpr (CharmHad == CharmHadAlice3::Lc) { + isSelPid(cand.nSigTrkPi2(), PidSels::TrkProng2, ptProng2, nSigmaTrkMax, ptPidTrkMin, ptPidTrkMax); + isSelPid(cand.nSigRichPi2(), PidSels::RichProng2, ptProng2, nSigmaRichMax, ptPidRichMin, ptPidRichMax); + isSelPid(cand.nSigInnTofPi2(), PidSels::InnTofProng2, ptProng2, nSigmaInnTofMax, ptPidInnTofMin, ptPidInnTofMax); + isSelPid(cand.nSigOutTofPi2(), PidSels::OutTofProng2, ptProng2, nSigmaOutTofMax, ptPidOutTofMin, ptPidOutTofMax); + } + + return; + } + + /// \brief function to apply Lc selections + /// \tparam CharmHad is charm hadron type + /// \tparam CandType is candidate type + /// \param cands are 3-prong candidates + template + void runSelect3Prong(CandType const& cands) + { + bool isSelMassHypo0{false}; + bool isSelMassHypo1{false}; + std::vector outputMl{-1.f, -1.f, -1.f}; + uint32_t pidMask = 0; + + // looping over 3-prong cands + for (const auto& cand : cands) { + outputMl = {-1.f, -1.f, -1.f}; + pidMask = 0; + + auto ptCand = cand.pt(); + int const ptBin = findBin(binsPt, ptCand); + if (ptBin == -1) { + candSelFlags(isSelMassHypo0, isSelMassHypo1, pidMask); + if (applyMl) { + candMlScores(outputMl[0], outputMl[1], outputMl[2]); + } + continue; + } + + // Here all cands pass the cut on the mass selection + bool const selMassHypo0 = selectionCandidateMass(ptBin, cand); + bool const selMassHypo1 = selectionCandidateMass(ptBin, cand); + if (!selMassHypo0 && !selMassHypo1) { + candSelFlags(isSelMassHypo0, isSelMassHypo1, pidMask); + if (applyMl) { + candMlScores(outputMl[0], outputMl[1], outputMl[2]); + } + continue; + } + if (activateQA) { + registry.fill(HIST("hSelections"), 2 + aod::SelectionStep::RecoSkims, ptCand); + } + + // Topological selection (TODO: track quality selection) + if (!selectionTopol(cand, ptCand)) { + candSelFlags(isSelMassHypo0, isSelMassHypo1, pidMask); + if (applyMl) { + candMlScores(outputMl[0], outputMl[1], outputMl[2]); + } + continue; + } + if (activateQA) { + registry.fill(HIST("hSelections"), 2 + aod::SelectionStep::RecoTopol, ptCand); + } + + // PID selection + configurePidMask(cand, pidMask); + if (pidMask == 0) { + candSelFlags(isSelMassHypo0, isSelMassHypo1, pidMask); + if (applyMl) { + candMlScores(outputMl[0], outputMl[1], outputMl[2]); + } + continue; + } + if (activateQA) { + registry.fill(HIST("hSelections"), 2 + aod::SelectionStep::RecoPID, ptCand); + } + + bool isSelectedMl = true; + // ML selections + if (applyMl) { + + std::vector inputFeaturesMassHypo0 = mlResponse.getInputFeatures(cand); + isSelectedMl = mlResponse.isSelectedMl(inputFeaturesMassHypo0, ptCand, outputMl); + candMlScores(outputMl[0], outputMl[1], outputMl[2]); + if (!isSelectedMl) { + candSelFlags(isSelMassHypo0, isSelMassHypo1, pidMask); + continue; + } + + if (activateQA) { + registry.fill(HIST("hSelections"), 2 + aod::SelectionStep::RecoMl, ptCand); + } + } + + candSelFlags(selMassHypo0, selMassHypo1, pidMask); + } + } + + /// \brief process function for cand selection + /// \param cands Lc cand table + void processLc(CandsLcWMcTruth const& cands) + { + runSelect3Prong(cands); + } + PROCESS_SWITCH(Alice3HfSelector3Prong, processLc, "Process 3 prong selection for Lc", true); +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +} diff --git a/ALICE3/TableProducer/alice3HfTreeCreator3Prong.cxx b/ALICE3/TableProducer/alice3HfTreeCreator3Prong.cxx new file mode 100644 index 00000000000..15bfcd8da17 --- /dev/null +++ b/ALICE3/TableProducer/alice3HfTreeCreator3Prong.cxx @@ -0,0 +1,517 @@ +// 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 alice3HfTreeCreator3Prong.cxx +/// \brief Writer of 3-prong candidates in the form of flat tables to be stored in TTrees. +/// Intended for debug, local optimization of analysis on small samples or ML training. +/// +/// \author Marcello Di Costanzo , Turin Polytechnic University and INFN Turin + +#include "ALICE3/DataModel/A3DecayFinderTables.h" +#include "ALICE3/DataModel/OTFPIDTrk.h" +#include "ALICE3/DataModel/OTFRICH.h" +#include "ALICE3/DataModel/OTFTOF.h" +#include "ALICE3/DataModel/RICH.h" +#include "ALICE3/Utils/utilsHfAlice3.h" +#include "Common/Core/RecoDecay.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +using namespace o2; +using namespace o2::analysis; +using namespace o2::framework; +using namespace o2::framework::expressions; + +namespace o2::aod +{ +namespace full +{ +DECLARE_SOA_COLUMN(RSecondaryVertex, rSecondaryVertex, float); //! Radius of secondary vertex (cm) +DECLARE_SOA_COLUMN(PtProng0, ptProng0, float); //! Transverse momentum of prong0 (GeV/c) +DECLARE_SOA_COLUMN(PProng0, pProng0, float); //! Momentum of prong0 (GeV/c) +DECLARE_SOA_COLUMN(ImpactParameterNormalised0, impactParameterNormalised0, float); //! Normalised impact parameter of prong0 +DECLARE_SOA_COLUMN(PtProng1, ptProng1, float); //! Transverse momentum of prong1 (GeV/c) +DECLARE_SOA_COLUMN(PProng1, pProng1, float); //! Momentum of prong1 (in GeV/c) +DECLARE_SOA_COLUMN(ImpactParameterNormalised1, impactParameterNormalised1, float); //! Normalised impact parameter of prong1 +DECLARE_SOA_COLUMN(PtProng2, ptProng2, float); //! Transverse momentum of prong2 (GeV/c) +DECLARE_SOA_COLUMN(PProng2, pProng2, float); //! Momentum of prong2 (GeV/c) +DECLARE_SOA_COLUMN(ImpactParameterNormalised2, impactParameterNormalised2, float); //! Normalised impact parameter of prong2 +DECLARE_SOA_COLUMN(M, m, float); //! Invariant mass of cand (GeV/c2) +DECLARE_SOA_COLUMN(Pt, pt, float); //! Transverse momentum of cand (GeV/c) +DECLARE_SOA_COLUMN(P, p, float); //! Momentum of cand (GeV/c) +DECLARE_SOA_COLUMN(Y, y, float); //! Rapidity of cand +DECLARE_SOA_COLUMN(Eta, eta, float); //! Pseudorapidity of cand +DECLARE_SOA_COLUMN(Phi, phi, float); //! Azimuth angle of cand +DECLARE_SOA_COLUMN(E, e, float); //! Energy of cand (GeV) +DECLARE_SOA_COLUMN(DecayLength, decayLength, float); //! Decay length of cand (cm) +DECLARE_SOA_COLUMN(DecayLengthXY, decayLengthXY, float); //! Transverse decay length of cand (cm) +DECLARE_SOA_COLUMN(DecayLengthNormalised, decayLengthNormalised, float); //! Normalised decay length of cand +DECLARE_SOA_COLUMN(DecayLengthXYNormalised, decayLengthXYNormalised, float); //! Normalised transverse decay length of cand +DECLARE_SOA_COLUMN(Cpa, cpa, float); //! Cosine pointing angle of cand +DECLARE_SOA_COLUMN(CpaXY, cpaXY, float); //! Cosine pointing angle of cand in transverse plane +// PID columns +// Prong 0 +DECLARE_SOA_COLUMN(NSigTrkPi0, nSigTrkPi0, float); //! +DECLARE_SOA_COLUMN(NSigRichPi0, nSigRichPi0, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPi0, nSigInnTofPi0, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPi0, nSigOutTofPi0, float); //! +DECLARE_SOA_COLUMN(NSigTrkKa0, nSigTrkKa0, float); //! +DECLARE_SOA_COLUMN(NSigRichKa0, nSigRichKa0, float); //! +DECLARE_SOA_COLUMN(NSigInnTofKa0, nSigInnTofKa0, float); //! +DECLARE_SOA_COLUMN(NSigOutTofKa0, nSigOutTofKa0, float); //! +DECLARE_SOA_COLUMN(NSigTrkPr0, nSigTrkPr0, float); //! +DECLARE_SOA_COLUMN(NSigRichPr0, nSigRichPr0, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPr0, nSigInnTofPr0, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPr0, nSigOutTofPr0, float); //! +// Prong 1 +DECLARE_SOA_COLUMN(NSigTrkPi1, nSigTrkPi1, float); //! +DECLARE_SOA_COLUMN(NSigRichPi1, nSigRichPi1, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPi1, nSigInnTofPi1, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPi1, nSigOutTofPi1, float); //! +DECLARE_SOA_COLUMN(NSigTrkKa1, nSigTrkKa1, float); //! +DECLARE_SOA_COLUMN(NSigRichKa1, nSigRichKa1, float); //! +DECLARE_SOA_COLUMN(NSigInnTofKa1, nSigInnTofKa1, float); //! +DECLARE_SOA_COLUMN(NSigOutTofKa1, nSigOutTofKa1, float); //! +DECLARE_SOA_COLUMN(NSigTrkPr1, nSigTrkPr1, float); //! +DECLARE_SOA_COLUMN(NSigRichPr1, nSigRichPr1, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPr1, nSigInnTofPr1, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPr1, nSigOutTofPr1, float); //! +// Prong 2 +DECLARE_SOA_COLUMN(NSigTrkPi2, nSigTrkPi2, float); //! +DECLARE_SOA_COLUMN(NSigRichPi2, nSigRichPi2, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPi2, nSigInnTofPi2, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPi2, nSigOutTofPi2, float); //! +DECLARE_SOA_COLUMN(NSigTrkKa2, nSigTrkKa2, float); //! +DECLARE_SOA_COLUMN(NSigRichKa2, nSigRichKa2, float); //! +DECLARE_SOA_COLUMN(NSigInnTofKa2, nSigInnTofKa2, float); //! +DECLARE_SOA_COLUMN(NSigOutTofKa2, nSigOutTofKa2, float); //! +DECLARE_SOA_COLUMN(NSigTrkPr2, nSigTrkPr2, float); //! +DECLARE_SOA_COLUMN(NSigRichPr2, nSigRichPr2, float); //! +DECLARE_SOA_COLUMN(NSigInnTofPr2, nSigInnTofPr2, float); //! +DECLARE_SOA_COLUMN(NSigOutTofPr2, nSigOutTofPr2, float); //! +// ML scores +DECLARE_SOA_COLUMN(MlScore0, mlScore0, float); //! ML score of the first configured index +DECLARE_SOA_COLUMN(MlScore1, mlScore1, float); //! ML score of the second configured index +DECLARE_SOA_COLUMN(MlScore2, mlScore2, float); //! ML score of the third configured index +} // namespace full + +// Topology tables +DECLARE_SOA_TABLE(Alice3CandVtxs, "AOD", "ALICE3CANDVTX", //! + collision::PosX, + collision::PosY, + collision::PosZ, + a3_hf_cand::XSecondaryVertex, + a3_hf_cand::YSecondaryVertex, + a3_hf_cand::ZSecondaryVertex, + full::RSecondaryVertex); + +DECLARE_SOA_TABLE(Alice3CandTopos, "AOD", "ALICE3CANDTOPO", //! + a3_hf_cand::Chi2PCA, + full::DecayLength, + a3_hf_cand::ErrorDecayLength, + full::DecayLengthXY, + a3_hf_cand::ErrorDecayLengthXY, + full::DecayLengthNormalised, + full::DecayLengthXYNormalised, + full::Cpa, + full::CpaXY); + +DECLARE_SOA_TABLE(Alice3DaugTopos, "AOD", "ALICE3DAUGTOPO", //! + full::ImpactParameterNormalised0, + full::ImpactParameterNormalised1, + full::ImpactParameterNormalised2, + a3_hf_cand::ImpactParameterY0, + a3_hf_cand::ImpactParameterY1, + a3_hf_cand::ImpactParameterY2, + a3_hf_cand::ErrorImpactParameterY0, + a3_hf_cand::ErrorImpactParameterY1, + a3_hf_cand::ErrorImpactParameterY2, + a3_hf_cand::ImpactParameterZ0, + a3_hf_cand::ImpactParameterZ1, + a3_hf_cand::ImpactParameterZ2, + a3_hf_cand::ErrorImpactParameterZ0, + a3_hf_cand::ErrorImpactParameterZ1, + a3_hf_cand::ErrorImpactParameterZ2); + +// ML tables +DECLARE_SOA_TABLE(Alice3PMls, "AOD", "ALICE3PML", + full::MlScore0, + full::MlScore1, + full::MlScore2) +// Kinematics information table +DECLARE_SOA_TABLE(Alice3Kine3Ps, "AOD", "ALICE3KINE3P", + full::PtProng0, + full::PtProng1, + full::PtProng2, + full::Pt, + full::P, + full::Eta, + full::Phi, + full::M, + full::Y, + full::E); + +DECLARE_SOA_TABLE(Alice3Cand3PGens, "AOD", "ALICE3CAND3PGEN", + full::Pt, + full::Eta, + full::Phi, + full::Y); +DECLARE_SOA_TABLE(Alice3DaugMoms, "AOD", "ALICE3DAUGMOM", //! + a3_hf_cand::PxProng0, a3_hf_cand::PyProng0, a3_hf_cand::PzProng0, full::PProng0, + a3_hf_cand::PxProng1, a3_hf_cand::PyProng1, a3_hf_cand::PzProng1, full::PProng1, + a3_hf_cand::PxProng2, a3_hf_cand::PyProng2, a3_hf_cand::PzProng2, full::PProng2); +// MC matching tables +DECLARE_SOA_TABLE(Alice3McRecs, "AOD", "ALICE3MCREC", + a3_mc_truth::FlagMcRec, + a3_mc_truth::OriginMcRec); + +DECLARE_SOA_TABLE(Alice3McGens, "AOD", "ALICE3MCGEN", + a3_mc_truth::FlagMcGen, + a3_mc_truth::OriginMcGen); +// PID tables +DECLARE_SOA_TABLE(Alice3PidPi0s, "AOD", "ALICE3PIDPI0", + full::NSigTrkPi0, + full::NSigRichPi0, + full::NSigInnTofPi0, + full::NSigOutTofPi0) +DECLARE_SOA_TABLE(Alice3PidPi1s, "AOD", "ALICE3PIDPI1", + full::NSigTrkPi1, + full::NSigRichPi1, + full::NSigInnTofPi1, + full::NSigOutTofPi1) +DECLARE_SOA_TABLE(Alice3PidPi2s, "AOD", "ALICE3PIDPI2", + full::NSigTrkPi2, + full::NSigRichPi2, + full::NSigInnTofPi2, + full::NSigOutTofPi2) +DECLARE_SOA_TABLE(Alice3PidKa0s, "AOD", "ALICE3PIDKA0", + full::NSigTrkKa0, + full::NSigRichKa0, + full::NSigInnTofKa0, + full::NSigOutTofKa0) +DECLARE_SOA_TABLE(Alice3PidKa1s, "AOD", "ALICE3PIDKA1", + full::NSigTrkKa1, + full::NSigRichKa1, + full::NSigInnTofKa1, + full::NSigOutTofKa1) +DECLARE_SOA_TABLE(Alice3PidKa2s, "AOD", "ALICE3PIDKA2", + full::NSigTrkKa2, + full::NSigRichKa2, + full::NSigInnTofKa2, + full::NSigOutTofKa2) +DECLARE_SOA_TABLE(Alice3PidPr0s, "AOD", "ALICE3PIDPR0", + full::NSigTrkPr0, + full::NSigRichPr0, + full::NSigInnTofPr0, + full::NSigOutTofPr0) +DECLARE_SOA_TABLE(Alice3PidPr1s, "AOD", "ALICE3PIDPR1", + full::NSigTrkPr1, + full::NSigRichPr1, + full::NSigInnTofPr1, + full::NSigOutTofPr1) +DECLARE_SOA_TABLE(Alice3PidPr2s, "AOD", "ALICE3PIDPR2", + full::NSigTrkPr2, + full::NSigRichPr2, + full::NSigInnTofPr2, + full::NSigOutTofPr2) + +} // namespace o2::aod + +/// Writes the full information in an output TTree +struct Alice3HfTreeCreator3Prong { + Produces rowCandVtxs; + Produces rowCandTopos; + Produces rowDaugTopos; + Produces rowCandMls; + Produces rowCandKine3Ps; + Produces rowCand3PGen; + Produces rowDaugMoms; + Produces rowCand3PMcMatchRec; + Produces rowCand3PMcMatchGen; + Produces rowPidPi0; + Produces rowPidPi1; + Produces rowPidPi2; + Produces rowPidKa0; + Produces rowPidKa1; + Produces rowPidKa2; + Produces rowPidPr0; + Produces rowPidPr1; + Produces rowPidPr2; + + // Configurables to fill tables + struct : o2::framework::ConfigurableGroup { + Configurable fillCandVtxInfo{"fillCandVtxInfo", false, "fill candidate vtx info"}; + Configurable fillCandTopoInfo{"fillCandTopoInfo", false, "fill candidate topology info"}; + Configurable fillDaugTopoInfo{"fillDaugTopoInfo", false, "fill daughter topology info"}; + Configurable fillMlScoreInfo{"fillMlScoreInfo", false, "fill ML scores info"}; + Configurable fillCandKineInfo{"fillCandKineInfo", false, "fill candidate kinematic info"}; + Configurable fillCandGenKineInfo{"fillCandGenKineInfo", false, "fill generated candidate kinematic info"}; + Configurable fillDaugKineInfo{"fillDaugKineInfo", false, "fill daughter kinematic info"}; + Configurable fillMcMatchRecoInfo{"fillMcMatchRecoInfo", false, "fill MC match reconstruction info"}; + Configurable fillMcMatchGenInfo{"fillMcMatchGenInfo", false, "fill MC match generation info"}; + Configurable fillPid{"fillPid", false, "fill PID info"}; + } fillTables; + // parameters for production of training samples + Configurable fillOnlySignal{"fillOnlySignal", true, "Flag to fill derived tables with signal"}; + Configurable fillOnlyBackground{"fillOnlyBackground", false, "Flag to fill derived tables with background"}; + Configurable downSampleFactor{"downSampleFactor", 1., "Fraction of cands to keep"}; + Configurable ptMaxForDownSample{"ptMaxForDownSample", 10., "Maximum pt for the application of the downsampling factor"}; + + HfHelperAlice3 hfHelper; + + using CandsLcRec = soa::Filtered>; + using CandsLcRecWMl = soa::Filtered>; + using CandsMcGen = soa::Join; + + Filter filterSelectCandidates = (aod::a3_hf_sel_3prong::isSelMassHypo0 == true || aod::a3_hf_sel_3prong::isSelMassHypo1 == true); + Filter filterSelectGenCands = nabs(aod::a3_mc_truth::flagMcGen) == static_cast(CharmHadAlice3::Lc); + + Partition recoLcCandSig = nabs(o2::aod::a3_mc_truth::flagMcRec) == static_cast(CharmHadAlice3::Lc); + Partition recoLcCandBkg = nabs(o2::aod::a3_mc_truth::flagMcRec) == 0; + Partition recoLcCandSigWMl = nabs(o2::aod::a3_mc_truth::flagMcRec) == static_cast(CharmHadAlice3::Lc); + Partition recoLcCandBkgWMl = nabs(o2::aod::a3_mc_truth::flagMcRec) == 0; + + void init(InitContext const&) + { + const std::array doprocess{doprocessLc, doprocessLcWMl}; + if ((std::accumulate(doprocess.begin(), doprocess.end(), 0)) != 1) { + LOGP(fatal, "no or more than one process function enabled! Please check your configuration!"); + } + } + + /// Reserve space in the output tables + /// \tparam TCand Type of candidate + /// \param nCands Number of candidates to reserve space for + /// \param cand Candidate to be used to check which PID tables to reserve + template + void reserveTables(size_t nCands, const TCand& cand) + { + if (fillTables.fillCandVtxInfo) + rowCandVtxs.reserve(nCands); + if (fillTables.fillCandTopoInfo) + rowCandTopos.reserve(nCands); + if (fillTables.fillDaugTopoInfo) + rowDaugTopos.reserve(nCands); + if (fillTables.fillMlScoreInfo) + rowCandMls.reserve(nCands); + if (fillTables.fillCandKineInfo) + rowCandKine3Ps.reserve(nCands); + if (fillTables.fillCandGenKineInfo) + rowCand3PGen.reserve(nCands); + if (fillTables.fillDaugKineInfo) + rowDaugMoms.reserve(nCands); + if (fillTables.fillMcMatchRecoInfo) + rowCand3PMcMatchRec.reserve(nCands); + if (fillTables.fillMcMatchGenInfo) + rowCand3PMcMatchGen.reserve(nCands); + // PID tables + if (fillTables.fillPid) { + if constexpr (requires { cand.nSigTrkPi0(); }) + rowPidPi0.reserve(nCands); + if constexpr (requires { cand.nSigTrkPi1(); }) + rowPidPi1.reserve(nCands); + if constexpr (requires { cand.nSigTrkPi2(); }) + rowPidPi2.reserve(nCands); + if constexpr (requires { cand.nSigTrkKa0(); }) + rowPidKa0.reserve(nCands); + if constexpr (requires { cand.nSigTrkKa1(); }) + rowPidKa1.reserve(nCands); + if constexpr (requires { cand.nSigTrkKa2(); }) + rowPidKa2.reserve(nCands); + if constexpr (requires { cand.nSigTrkPr0(); }) + rowPidPr0.reserve(nCands); + if constexpr (requires { cand.nSigTrkPr1(); }) + rowPidPr1.reserve(nCands); + if constexpr (requires { cand.nSigTrkPr2(); }) + rowPidPr2.reserve(nCands); + } + } + + /// Fill reconstructed candidate tables + /// \tparam CharmHadAlice3: charm hadron type + /// \tparam IsSwapMassHypo: whether to swap mass hypothesis or not + /// \tparam T: candidate type + /// \param cand: candidate to be used to fill the tables + template + void fillRecoTables(const T& cand) + { + if (fillTables.fillCandVtxInfo) { + rowCandVtxs( + cand.posX(), + cand.posY(), + cand.posZ(), + cand.xSecondaryVertex(), + cand.ySecondaryVertex(), + cand.zSecondaryVertex(), + cand.rSecondaryVertex()); + } + if (fillTables.fillCandTopoInfo) { + rowCandTopos( + cand.chi2PCA(), + cand.decayLength(), + cand.errorDecayLength(), + cand.decayLengthXY(), + cand.errorDecayLengthXY(), + cand.decayLengthNormalised(), + cand.decayLengthXYNormalised(), + cand.cpa(), + cand.cpaXY()); + } + if (fillTables.fillDaugTopoInfo) { + rowDaugTopos( + cand.impactParameterNormalised0(), + cand.impactParameterNormalised1(), + cand.impactParameterNormalised2(), + cand.impactParameterY0(), + cand.impactParameterY1(), + cand.impactParameterY2(), + cand.errorImpactParameterY0(), + cand.errorImpactParameterY1(), + cand.errorImpactParameterY2(), + cand.impactParameterZ0(), + cand.impactParameterZ1(), + cand.impactParameterZ2(), + cand.errorImpactParameterZ0(), + cand.errorImpactParameterZ1(), + cand.errorImpactParameterZ2()); + } + if constexpr (requires { cand.mlScore0(); }) { + if (fillTables.fillMlScoreInfo) { + rowCandMls(cand.mlScore0(), cand.mlScore1(), cand.mlScore2()); + } + } + if (fillTables.fillCandKineInfo) { + rowCandKine3Ps( + cand.ptProng0(), cand.ptProng1(), cand.ptProng2(), + cand.pt(), RecoDecay::p(cand.px(), cand.py(), cand.pz()), cand.eta(), cand.phi(), + hfHelper.getCandMass(cand), hfHelper.getCandY(cand), hfHelper.getCandEnergy(cand)); + } + if (fillTables.fillDaugKineInfo) { + rowDaugMoms( + cand.pxProng0(), cand.pyProng0(), cand.pzProng0(), RecoDecay::p(cand.pxProng0(), cand.pyProng0(), cand.pzProng0()), + cand.pxProng1(), cand.pyProng1(), cand.pzProng1(), RecoDecay::p(cand.pxProng1(), cand.pyProng1(), cand.pzProng1()), + cand.pxProng2(), cand.pyProng2(), cand.pzProng2(), RecoDecay::p(cand.pxProng2(), cand.pyProng2(), cand.pzProng2())); + } + if (fillTables.fillMcMatchRecoInfo) + rowCand3PMcMatchRec(cand.flagMcRec(), cand.originMcRec()); + + // Fill PID tables + if (fillTables.fillPid) { + if constexpr (requires { cand.nSigTrkPi0(); }) + rowPidPi0(cand.nSigTrkPi0(), cand.nSigRichPi0(), cand.nSigInnTofPi0(), cand.nSigOutTofPi0()); + if constexpr (requires { cand.nSigTrkPi1(); }) + rowPidPi1(cand.nSigTrkPi1(), cand.nSigRichPi1(), cand.nSigInnTofPi1(), cand.nSigOutTofPi1()); + if constexpr (requires { cand.nSigTrkPi2(); }) + rowPidPi2(cand.nSigTrkPi2(), cand.nSigRichPi2(), cand.nSigInnTofPi2(), cand.nSigOutTofPi2()); + if constexpr (requires { cand.nSigTrkKa0(); }) + rowPidKa0(cand.nSigTrkKa0(), cand.nSigRichKa0(), cand.nSigInnTofKa0(), cand.nSigOutTofKa0()); + if constexpr (requires { cand.nSigTrkKa1(); }) + rowPidKa1(cand.nSigTrkKa1(), cand.nSigRichKa1(), cand.nSigInnTofKa1(), cand.nSigOutTofKa1()); + if constexpr (requires { cand.nSigTrkKa2(); }) + rowPidKa2(cand.nSigTrkKa2(), cand.nSigRichKa2(), cand.nSigInnTofKa2(), cand.nSigOutTofKa2()); + if constexpr (requires { cand.nSigTrkPr0(); }) + rowPidPr0(cand.nSigTrkPr0(), cand.nSigRichPr0(), cand.nSigInnTofPr0(), cand.nSigOutTofPr0()); + if constexpr (requires { cand.nSigTrkPr1(); }) + rowPidPr1(cand.nSigTrkPr1(), cand.nSigRichPr1(), cand.nSigInnTofPr1(), cand.nSigOutTofPr1()); + if constexpr (requires { cand.nSigTrkPr2(); }) + rowPidPr2(cand.nSigTrkPr2(), cand.nSigRichPr2(), cand.nSigInnTofPr2(), cand.nSigOutTofPr2()); + } + } + + /// Function to fill generated tables + /// \tparam CharmHad Type of 3prong particle + /// \tparam T Type of generated candidates collection + /// \param parts Generated candidates collection + template + void fillGenTables(const T& parts) + { + for (const auto& part : parts) { + if (fillTables.fillCandGenKineInfo) { + rowCand3PGen(part.pt(), part.eta(), part.phi(), hfHelper.getCandY(part)); + } + if (fillTables.fillMcMatchGenInfo) { + rowCand3PMcMatchGen(part.flagMcGen(), part.originMcGen()); + } + } + } + + /// Function to fill both reco and gen tables + /// from any candidate collection + /// \tparam CharmHad Type of 3prong particle + /// \tparam TCandsRec Type of reconstructed candidates collection + /// \tparam TCandsGen Type of generated candidates collection + /// \param candsRec Reconstructed candidates collection + /// \param candsGen Generated candidates collection + template + void fillRecoGenTables(const TCandsRec& candsRec, + const TCandsGen& candsGen) + { + reserveTables(candsRec.size(), *candsRec.begin()); + for (const auto& cand : candsRec) { + if (downSampleFactor < 1.) { + float const pseudoRndm = cand.ptProng0() * 1000. - static_cast(cand.ptProng0() * 1000); + if (cand.pt() < ptMaxForDownSample && pseudoRndm >= downSampleFactor) { + continue; + } + } + if (cand.isSelMassHypo0()) { + fillRecoTables(cand); + } + if (cand.isSelMassHypo1()) { + fillRecoTables(cand); + } + } + fillGenTables(candsGen); + } + + void processLc(CandsLcRec const& selCandsLcRec, + CandsMcGen const& parts, + aod::McCollisions const&) + { + if (fillOnlySignal) { + fillRecoGenTables(recoLcCandSig, parts); + } else if (fillOnlyBackground) { + fillRecoGenTables(recoLcCandBkg, parts); + } else { + fillRecoGenTables(selCandsLcRec, parts); + } + } + PROCESS_SWITCH(Alice3HfTreeCreator3Prong, processLc, "Process Lc", true); + + void processLcWMl(CandsLcRecWMl const& selCandsLcRec, + CandsMcGen const& parts, + aod::McCollisions const&) + { + if (fillOnlySignal) { + fillRecoGenTables(recoLcCandSig, parts); + } else if (fillOnlyBackground) { + fillRecoGenTables(recoLcCandBkg, parts); + } else { + fillRecoGenTables(selCandsLcRec, parts); + } + } + PROCESS_SWITCH(Alice3HfTreeCreator3Prong, processLcWMl, "Process Lc with ML", false); +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +} diff --git a/ALICE3/Tasks/CMakeLists.txt b/ALICE3/Tasks/CMakeLists.txt index 42fb53a0a25..a0ed8795f1f 100644 --- a/ALICE3/Tasks/CMakeLists.txt +++ b/ALICE3/Tasks/CMakeLists.txt @@ -64,6 +64,11 @@ o2physics_add_dpl_workflow(alice3-multicharm PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore COMPONENT_NAME Analysis) +o2physics_add_dpl_workflow(alice3-hf-task-3prong + SOURCES alice3HfTask3Prong.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::SGCutParHolder O2Physics::MLCore + COMPONENT_NAME Analysis) + o2physics_add_dpl_workflow(alice3-efficiency SOURCES alice3Efficiency.cxx PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore diff --git a/ALICE3/Tasks/alice3HfTask3Prong.cxx b/ALICE3/Tasks/alice3HfTask3Prong.cxx new file mode 100644 index 00000000000..fa0e2867503 --- /dev/null +++ b/ALICE3/Tasks/alice3HfTask3Prong.cxx @@ -0,0 +1,367 @@ +// 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 hfTask3Prong.cxx +/// \brief 3-prong candidates analysis task for ALICE 3 simulation studies +/// \author Marcello Di Costanzo , Polytechnic University of Turin and INFN Turin + +#include "ALICE3/DataModel/A3DecayFinderTables.h" +#include "ALICE3/DataModel/OTFPIDTrk.h" +#include "ALICE3/DataModel/OTFRICH.h" +#include "ALICE3/DataModel/OTFTOF.h" +#include "ALICE3/DataModel/RICH.h" +#include "ALICE3/Utils/utilsHfAlice3.h" +#include "ALICE3/Utils/utilsSelectionsAlice3.h" +#include "Common/Core/RecoDecay.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +using namespace o2; +using namespace o2::analysis; +using namespace o2::framework; +using namespace o2::framework::expressions; + +/// Λc± → p± K∓ π± analysis task +struct Alice3HfTask3Prong { + Configurable yCandGenMax{"yCandGenMax", 0.8, "max. gen particle rapidity"}; + Configurable yCandRecoMax{"yCandRecoMax", 0.8, "max. cand. rapidity"}; + Configurable> binsPt{"binsPt", std::vector{hf_cuts_3prongs_alice3::vecBinsPt}, "pT bin limits"}; + Configurable fillThn{"fillThn", false, "fill Thn"}; + Configurable ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; + Configurable ccdbPathGrp{"ccdbPathGrp", "GLO/GRP/GRP", "Path of the grp file (Run 2)"}; + Configurable ccdbPathGrpMag{"ccdbPathGrpMag", "GLO/Config/GRPMagField", "CCDB path of the GRPMagField object (Run 3)"}; + + HfHelperAlice3 hfHelper; + SliceCache cache; + Service ccdb; + + int selectedPdg{-1}; + + using Cands3PReco = soa::Filtered>; + using Cands3PRecoWMl = soa::Filtered>; + using Cands3PGen = soa::Join; + + Filter filterSelectCandidates = (aod::a3_hf_sel_3prong::isSelMassHypo0 == true || aod::a3_hf_sel_3prong::isSelMassHypo1 == true); + + Partition candsGenLcs = nabs(aod::a3_mc_truth::flagMcGen) == static_cast(CharmHadAlice3::Lc); + + ConfigurableAxis thnConfigAxisPt{"thnConfigAxisPt", {72, 0, 36}, ""}; + ConfigurableAxis thnConfigAxisMass{"thnConfigAxisMass", {300, 1.98, 2.58}, ""}; + ConfigurableAxis thnConfigAxisBdtScoreBkg{"thnConfigAxisBdtScoreBkg", {1000, 0., 1.}, ""}; + ConfigurableAxis thnConfigAxisBdtScoreSignal{"thnConfigAxisBdtScoreSignal", {100, 0., 1.}, ""}; + ConfigurableAxis thnConfigAxisCanType{"thnConfigAxisCanType", {5, 0., 5.}, ""}; + ConfigurableAxis thnAxisRapidity{"thnAxisRapidity", {20, -1, 1}, "Cand. rapidity bins"}; + ConfigurableAxis thnConfigAxisGenPtB{"thnConfigAxisGenPtB", {1000, 0, 100}, "Gen Pt B"}; + + HistogramRegistry registry{"registry", {}}; + + // Names of folders and suffixes for MC signal histograms + constexpr static std::string_view SignalFolders[] = {"signal", "prompt", "nonprompt"}; + constexpr static std::string_view SignalSuffixes[] = {"", "Prompt", "NonPrompt"}; + + enum SignalClasses : int { + Signal = 0, + Prompt, + NonPrompt + }; + + void init(InitContext&) + { + const std::array doprocess{doprocessLc, doprocessLcWMl}; + if ((std::accumulate(doprocess.begin(), doprocess.end(), 0)) != 1) { + LOGP(fatal, "no or more than one process function enabled! Please check your configuration!"); + } + + if (doprocessLc || doprocessLcWMl) { + selectedPdg = CharmHadAlice3::Lc; + } + + auto addHistogramsRec = [&](const std::string& histoName, const std::string& xAxisTitle, const std::string& yAxisTitle, const HistogramConfigSpec& configSpec) { + registry.add(("MC/rec/signal/" + histoName + "RecSig").c_str(), ("3-prong cands (matched);" + xAxisTitle + ";" + yAxisTitle).c_str(), configSpec); + registry.add(("MC/rec/prompt/" + histoName + "RecSigPrompt").c_str(), ("3-prong cands (matched, prompt);" + xAxisTitle + ";" + yAxisTitle).c_str(), configSpec); + registry.add(("MC/rec/nonprompt/" + histoName + "RecSigNonPrompt").c_str(), ("3-prong cands (matched, non-prompt);" + xAxisTitle + ";" + yAxisTitle).c_str(), configSpec); + }; + + auto addHistogramsGen = [&](const std::string& histoName, const std::string& xAxisTitle, const std::string& yAxisTitle, const HistogramConfigSpec& configSpec) { + registry.add(("MC/gen/signal/" + histoName + "Gen").c_str(), ("MC particles (matched);" + xAxisTitle + ";" + yAxisTitle).c_str(), configSpec); + registry.add(("MC/gen/prompt/" + histoName + "GenPrompt").c_str(), ("MC particles (matched, prompt);" + xAxisTitle + ";" + yAxisTitle).c_str(), configSpec); + registry.add(("MC/gen/nonprompt/" + histoName + "GenNonPrompt").c_str(), ("MC particles (matched, non-prompt);" + xAxisTitle + ";" + yAxisTitle).c_str(), configSpec); + }; + + auto vbins = (std::vector)binsPt; + + /// Reconstructed Histograms + addHistogramsRec("hMass", "inv. mass (p K #pi) (GeV/#it{c}^{2})", "", {HistType::kTH1F, {{600, 1.98, 2.58}}}); + addHistogramsRec("hPt", "#it{p}_{T}^{rec.} (GeV/#it{c})", "entries", {HistType::kTH1F, {{360, 0., 36.}}}); + addHistogramsRec("hPhi", "#it{#Phi}", "entries", {HistType::kTH1F, {{100, 0., 6.3}}}); + addHistogramsRec("hPtProng0", "prong 0 #it{p}_{T} (GeV/#it{c})", "entries", {HistType::kTH1F, {{360, 0., 36.}}}); + addHistogramsRec("hPtProng1", "prong 1 #it{p}_{T} (GeV/#it{c})", "entries", {HistType::kTH1F, {{360, 0., 36.}}}); + addHistogramsRec("hPtProng2", "prong 2 #it{p}_{T} (GeV/#it{c})", "entries", {HistType::kTH1F, {{360, 0., 36.}}}); + addHistogramsRec("hd0Prong0", "prong 0 DCAxy to prim. vertex (cm)", "entries", {HistType::kTH1F, {{600, -0.4, 0.4}}}); + addHistogramsRec("hd0Prong1", "prong 1 DCAxy to prim. vertex (cm)", "entries", {HistType::kTH1F, {{600, -0.4, 0.4}}}); + addHistogramsRec("hd0Prong2", "prong 2 DCAxy to prim. vertex (cm)", "entries", {HistType::kTH1F, {{600, -0.4, 0.4}}}); + addHistogramsRec("hDecLength", "decay length (cm)", "entries", {HistType::kTH1F, {{400, 0., 1.}}}); + addHistogramsRec("hDecLengthxy", "decay length xy (cm)", "entries", {HistType::kTH1F, {{400, 0., 1.}}}); + addHistogramsRec("hCPA", "cosine of pointing angle", "entries", {HistType::kTH1F, {{110, -1.1, 1.1}}}); + addHistogramsRec("hCPAxy", "cosine of pointing angle xy", "entries", {HistType::kTH1F, {{110, -1.1, 1.1}}}); + addHistogramsRec("hDca2", "prong Chi2PCA to sec. vertex (cm)", "entries", {HistType::kTH1F, {{400, 0., 20.}}}); + addHistogramsRec("hEta", "#it{#eta}", "entries", {HistType::kTH1F, {{100, -2., 2.}}}); + addHistogramsRec("hMassVsPt", "inv. mass (p K #pi) (GeV/#it{c}^{2})", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{600, 1.98, 2.58}, {vbins}}}); + addHistogramsRec("hd0VsPtProng0", "prong 0 DCAxy to prim. vertex (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{600, -0.4, 0.4}, {vbins}}}); + addHistogramsRec("hd0VsPtProng1", "prong 1 DCAxy to prim. vertex (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{600, -0.4, 0.4}, {vbins}}}); + addHistogramsRec("hd0VsPtProng2", "prong 2 DCAxy to prim. vertex (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{600, -0.4, 0.4}, {vbins}}}); + addHistogramsRec("hDecLengthVsPt", "decay length (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{400, 0., 1.}, {vbins}}}); + addHistogramsRec("hDecLengthxyVsPt", "decay length xy (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{400, 0., 1.}, {vbins}}}); + addHistogramsRec("hCPAVsPt", "cosine of pointing angle", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{110, -1.1, 1.1}, {vbins}}}); + addHistogramsRec("hCPAxyVsPt", "cosine of pointing angle xy", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{110, -1.1, 1.1}, {vbins}}}); + addHistogramsRec("hDca2VsPt", "prong Chi2PCA to sec. vertex (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{400, 0., 20.}, {vbins}}}); + addHistogramsRec("hEtaVsPt", "candidate #it{#eta}", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{100, -2., 2.}, {vbins}}}); + addHistogramsRec("hPhiVsPt", "candidate #it{#Phi}", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{100, 0., 6.3}, {vbins}}}); + addHistogramsRec("hImpParErrProng0VsPt", "prong 0 impact parameter error (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{100, -1., 1.}, {vbins}}}); + addHistogramsRec("hImpParErrProng1VsPt", "prong 1 impact parameter error (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{100, -1., 1.}, {vbins}}}); + addHistogramsRec("hImpParErrProng2VsPt", "prong 2 impact parameter error (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{100, -1., 1.}, {vbins}}}); + addHistogramsRec("hDecLenErrVsPt", "decay length error (cm)", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{100, 0., 1.}, {vbins}}}); + + /// Generated Histograms + addHistogramsGen("hPt", "#it{p}_{T}^{gen.} (GeV/#it{c})", "entries", {HistType::kTH1F, {{360, 0., 36.}}}); + addHistogramsGen("hEta", "#it{#eta}", "entries", {HistType::kTH1F, {{100, -2., 2.}}}); + addHistogramsGen("hPhi", "#it{#Phi}", "entries", {HistType::kTH1F, {{100, 0., 6.3}}}); + addHistogramsGen("hY", "#it{y}", "entries", {HistType::kTH1F, {{100, -2., 2.}}}); + addHistogramsGen("hEtaVsPt", "#it{#eta}", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{100, -2., 2.}, {vbins}}}); + addHistogramsGen("hYVsPt", "#it{y}", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{100, -2., 2.}, {vbins}}}); + addHistogramsGen("hPhiVsPt", "#it{#Phi}", "#it{p}_{T} (GeV/#it{c})", {HistType::kTH2F, {{100, 0., 6.3}, {vbins}}}); + + /// selection status + registry.add("hSelectionStatus", "3-prong cands;selection status;entries", {HistType::kTH2F, {{5, -0.5, 4.5}, {vbins, "#it{p}_{T} (GeV/#it{c})"}}}); + + if (fillThn) { + const AxisSpec thnAxisMass{thnConfigAxisMass, "inv. mass (p K #pi) (GeV/#it{c}^{2})"}; + const AxisSpec thnAxisPt{thnConfigAxisPt, "#it{p}_{T}(#Lambda_{c}^{+}) (GeV/#it{c})"}; + const AxisSpec thnAxisScoreBkg{thnConfigAxisBdtScoreBkg, "BDT bkg score"}; + const AxisSpec thnAxisScorePrompt{thnConfigAxisBdtScoreSignal, "BDT prompt score"}; + const AxisSpec thnAxisScoreNonPrompt{thnConfigAxisBdtScoreSignal, "BDT non-prompt score"}; + const AxisSpec thnAxisCanType{thnConfigAxisCanType, "candidates type"}; + const AxisSpec thnAxisY{thnAxisRapidity, "rapidity"}; + const AxisSpec thnAxisPtB{thnConfigAxisGenPtB, "#it{p}_{T}^{B} (GeV/#it{c})"}; + + std::vector axesWithBdt = {thnAxisMass, thnAxisPt, thnAxisScoreBkg, thnAxisScorePrompt, thnAxisScoreNonPrompt, thnAxisPtB, thnAxisCanType}; + registry.add("hSparseRec", "Thn for reco cands", HistType::kTHnSparseF, axesWithBdt); + std::vector axesGen = {thnAxisPt, thnAxisY, thnAxisPtB, thnAxisCanType}; + registry.add("hSparseGen", "Thn for gen cands", HistType::kTHnSparseF, axesGen); + } + + ccdb->setURL(ccdbUrl); + ccdb->setCaching(true); + ccdb->setLocalObjectValidityChecking(); + } + + /// Helper function for filling MC reconstructed histograms for prompt, nonpromt and common (signal) + /// \param candidate is a reconstructed candidate + /// \tparam SignalType is an enum defining which histogram in which folder (signal, prompt or nonpromt) to fill + template + void fillHistogramsRecSig(CandidateType const& candidate, float mass) + { + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hMassRecSig") + HIST(SignalSuffixes[SignalType]), mass); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hMassVsPtRecSig") + HIST(SignalSuffixes[SignalType]), mass, candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hPtProng0RecSig") + HIST(SignalSuffixes[SignalType]), candidate.ptProng0()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hPtProng1RecSig") + HIST(SignalSuffixes[SignalType]), candidate.ptProng1()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hPtProng2RecSig") + HIST(SignalSuffixes[SignalType]), candidate.ptProng2()); + + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hd0Prong0RecSig") + HIST(SignalSuffixes[SignalType]), candidate.impactParameterY0()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hd0Prong1RecSig") + HIST(SignalSuffixes[SignalType]), candidate.impactParameterY1()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hd0Prong2RecSig") + HIST(SignalSuffixes[SignalType]), candidate.impactParameterY2()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hd0VsPtProng0RecSig") + HIST(SignalSuffixes[SignalType]), candidate.impactParameterY0(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hd0VsPtProng1RecSig") + HIST(SignalSuffixes[SignalType]), candidate.impactParameterY1(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hd0VsPtProng2RecSig") + HIST(SignalSuffixes[SignalType]), candidate.impactParameterY2(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hDecLengthRecSig") + HIST(SignalSuffixes[SignalType]), candidate.decayLength()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hDecLengthVsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.decayLength(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hDecLengthxyRecSig") + HIST(SignalSuffixes[SignalType]), candidate.decayLengthXY()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hDecLengthxyVsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.decayLengthXY(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hCPARecSig") + HIST(SignalSuffixes[SignalType]), candidate.cpa()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hCPAVsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.cpa(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hCPAxyRecSig") + HIST(SignalSuffixes[SignalType]), candidate.cpaXY()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hCPAxyVsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.cpaXY(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hDca2RecSig") + HIST(SignalSuffixes[SignalType]), candidate.chi2PCA()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hDca2VsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.chi2PCA(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hEtaRecSig") + HIST(SignalSuffixes[SignalType]), candidate.eta()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hEtaVsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.eta(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hPhiRecSig") + HIST(SignalSuffixes[SignalType]), candidate.phi()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hPhiVsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.phi(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hImpParErrProng0VsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.errorImpactParameterY0(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hImpParErrProng1VsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.errorImpactParameterY1(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hImpParErrProng2VsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.errorImpactParameterY2(), candidate.pt()); + registry.fill(HIST("MC/rec/") + HIST(SignalFolders[SignalType]) + HIST("/hDecLenErrVsPtRecSig") + HIST(SignalSuffixes[SignalType]), candidate.errorDecayLength(), candidate.pt()); + } + + /// Fill MC histograms at reconstruction level + /// \tparam CharmHad is the charm hadron species + /// \tparam SaveMl indicates whether ML scores are saved in the THnSparse + /// \tparam CandsRec is the type of the reconstructed candidates collection + /// \param candidates is the collection of reconstructed candidates + template + void fillHistosMcRec(CandsRec const& candidates) + { + for (const auto& candidate : candidates) { + /// rapidity selection + if (yCandRecoMax >= 0. && std::abs(hfHelper.getCandY(candidate)) > yCandRecoMax) { + continue; + } + + if (candidate.flagMcRec() != 0) { + // Get the corresponding MC particle. + + const auto pt = candidate.pt(); + const auto originType = candidate.originMcRec(); + + if (fillThn) { + if (candidate.isSelMassHypo0()) { + registry.fill(HIST("hSelectionStatus"), 0., pt); + double mass = hfHelper.getCandMass(candidate); + /// Fill histograms + fillHistogramsRecSig(candidate, mass); + if (originType == RecoDecay::OriginType::Prompt) { + fillHistogramsRecSig(candidate, mass); + } else if (originType == RecoDecay::OriginType::NonPrompt) { + fillHistogramsRecSig(candidate, mass); + } + std::vector valuesToFill{mass, pt}; + if constexpr (SaveMl) { + LOGP(fatal, "Trying to access ML scores, but SaveMl is false!"); + valuesToFill.push_back(candidate.mlScore0()); + valuesToFill.push_back(candidate.mlScore1()); + valuesToFill.push_back(candidate.mlScore2()); + } + valuesToFill.push_back(static_cast(originType)); + registry.get(HIST("hSparseRec"))->Fill(valuesToFill.data()); + } + if (candidate.isSelMassHypo1()) { + registry.fill(HIST("hSelectionStatus"), 1., pt); + double mass = hfHelper.getCandMass(candidate); + /// Fill histograms + fillHistogramsRecSig(candidate, mass); + if (originType == RecoDecay::OriginType::Prompt) { + fillHistogramsRecSig(candidate, mass); + } else if (originType == RecoDecay::OriginType::NonPrompt) { + fillHistogramsRecSig(candidate, mass); + } + std::vector valuesToFill{mass, pt}; + if constexpr (SaveMl) { + LOGP(fatal, "Trying to access ML scores, but SaveMl is false!"); + valuesToFill.push_back(candidate.mlScore0()); + valuesToFill.push_back(candidate.mlScore1()); + valuesToFill.push_back(candidate.mlScore2()); + } + valuesToFill.push_back(static_cast(originType)); + registry.get(HIST("hSparseRec"))->Fill(valuesToFill.data()); + } + } + } + } + } + + /// Helper function for filling MC generated histograms for prompt, nonpromt and common (signal) + /// \tparam CharmHad is the charm hadron species + /// \tparam SignalType is an enum defining which histogram in which folder (signal, prompt or nonpromt) to fill + /// \tparam ParticleType is the type of the generated particle + /// \param particle is a generated particle + template + void fillHistogramsGen(ParticleType const& particle) + { + LOG(debug) << "Filling generated histograms for signal type " << SignalType; + registry.fill(HIST("MC/gen/") + HIST(SignalFolders[SignalType]) + HIST("/hPtGen") + HIST(SignalSuffixes[SignalType]), particle.pt()); + registry.fill(HIST("MC/gen/") + HIST(SignalFolders[SignalType]) + HIST("/hEtaGen") + HIST(SignalSuffixes[SignalType]), particle.eta()); + registry.fill(HIST("MC/gen/") + HIST(SignalFolders[SignalType]) + HIST("/hYGen") + HIST(SignalSuffixes[SignalType]), hfHelper.getCandY(particle)); + registry.fill(HIST("MC/gen/") + HIST(SignalFolders[SignalType]) + HIST("/hPhiGen") + HIST(SignalSuffixes[SignalType]), particle.phi()); + registry.fill(HIST("MC/gen/") + HIST(SignalFolders[SignalType]) + HIST("/hEtaVsPtGen") + HIST(SignalSuffixes[SignalType]), particle.eta(), particle.pt()); + registry.fill(HIST("MC/gen/") + HIST(SignalFolders[SignalType]) + HIST("/hYVsPtGen") + HIST(SignalSuffixes[SignalType]), hfHelper.getCandY(particle), particle.pt()); + registry.fill(HIST("MC/gen/") + HIST(SignalFolders[SignalType]) + HIST("/hPhiVsPtGen") + HIST(SignalSuffixes[SignalType]), particle.phi(), particle.pt()); + } + + /// Fill MC histograms at generated level + /// \tparam CharmHad is the charm hadron species + /// \tparam CandsGen is the type of the generated candidates collection + /// \param mcParticles is the collection of generated particles + template + void fillHistosMcGen(CandsGen const& mcParticles) + { + // MC gen. + for (const auto& particle : mcParticles) { + if (std::abs(particle.flagMcGen()) == selectedPdg) { + double yGen = hfHelper.getCandY(particle); + if (yCandGenMax >= 0. && std::abs(yGen) > yCandGenMax) { + continue; + } + const auto ptGen = particle.pt(); + const auto originType = particle.originMcGen(); + + fillHistogramsGen(particle); + + float ptGenB = -1.f; + if (originType == RecoDecay::OriginType::Prompt) { + fillHistogramsGen(particle); + } else if (particle.originMcGen() == RecoDecay::OriginType::NonPrompt) { + ptGenB = particle.bHadMotherPtGen(); + fillHistogramsGen(particle); + } + + if (fillThn) { + std::vector valuesToFill{ptGen, yGen, ptGenB, static_cast(originType)}; + registry.get(HIST("hSparseGen"))->Fill(valuesToFill.data()); + } + } + } + } + + void processLc(Cands3PReco const& candsLc, + Cands3PGen const&) + { + fillHistosMcRec(candsLc); + fillHistosMcGen(candsGenLcs); + } + PROCESS_SWITCH(Alice3HfTask3Prong, processLc, "Process Lc w/o ML sels", true); + + void processLcWMl(Cands3PRecoWMl const& candsLcWMl, + Cands3PGen const&) + { + fillHistosMcRec(candsLcWMl); + fillHistosMcGen(candsGenLcs); + } + PROCESS_SWITCH(Alice3HfTask3Prong, processLcWMl, "Process Lc with ML sels", false); +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +} diff --git a/ALICE3/Utils/utilsHfAlice3.h b/ALICE3/Utils/utilsHfAlice3.h new file mode 100644 index 00000000000..4f58316369b --- /dev/null +++ b/ALICE3/Utils/utilsHfAlice3.h @@ -0,0 +1,97 @@ +// 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 HfHelperAlice3.h +/// \brief Class with helper functions for HF analyses +/// +/// \author Marcello Di Costanzo , Polytechnic University of Turin and INFN Turin + +#ifndef ALICE3_UTILS_UTILSHFALICE3_H_ +#define ALICE3_UTILS_UTILSHFALICE3_H_ + +#include "PWGHF/Core/HfHelper.h" + +namespace o2::analysis +{ + +enum CharmHadAlice3 { Lc = 1 }; + +} // namespace o2::analysis + +class HfHelperAlice3 +{ + public: + /// Default constructor + HfHelperAlice3() = default; + + /// Default destructor + ~HfHelperAlice3() = default; + + /// Get candidate mass (ALICE3 HF data model) + /// \tparam TCand candidate type + /// \param cand candidate + /// \return candidate mass + template + static double getCandMass(const TCand& cand) + { + switch (CharmHad) { + case o2::analysis::CharmHadAlice3::Lc: + return SwapHypo ? HfHelper::invMassLcToPiKP(cand) : HfHelper::invMassLcToPKPi(cand); + default: + LOG(fatal) << "Unsupported charm hadron type"; + return -1.; + } + } + + /// Get candidate energy (ALICE3 HF data model) + /// \tparam TCand candidate type + /// \param cand candidate + /// \return candidate energy + template + static double getCandEnergy(const TCand& cand) + { + switch (CharmHad) { + case o2::analysis::CharmHadAlice3::Lc: + return HfHelper::eLc(cand); + default: + LOG(fatal) << "Unsupported charm hadron type"; + return -1.; + } + } + + /// Get candidate rapidity (ALICE3 HF data model) + /// \tparam TCand candidate type + /// \param cand candidate + /// \return candidate rapidity + template + static double getCandY(const TCand& cand) + { + if constexpr (requires { cand.flagMcRec(); }) { + switch (CharmHad) { + case o2::analysis::CharmHadAlice3::Lc: + return HfHelper::yLc(cand); + default: + LOG(fatal) << "Unsupported charm hadron type"; + return -1.; + } + } else { + switch (CharmHad) { + case o2::analysis::CharmHadAlice3::Lc: + return RecoDecay::y(cand.pVector(), o2::constants::physics::MassLambdaCPlus); + default: + LOG(fatal) << "Unsupported charm hadron type"; + return -1.; + } + } + } +}; + +#endif // ALICE3_UTILS_UTILSHFALICE3_H_ diff --git a/ALICE3/Utils/utilsSelectionsAlice3.h b/ALICE3/Utils/utilsSelectionsAlice3.h new file mode 100644 index 00000000000..585b22e5295 --- /dev/null +++ b/ALICE3/Utils/utilsSelectionsAlice3.h @@ -0,0 +1,76 @@ +// 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 utilsSelectionsAlice3s.h +/// \brief Default pT bins and cut arrays for selections in ALICE3 performance analysis tasks +/// +/// \author Marcello Di Costanzo , Polytechnic University of Turin and INFN Turin + +#ifndef ALICE3_UTILS_UTILSSELECTIONSALICE3_H_ +#define ALICE3_UTILS_UTILSSELECTIONSALICE3_H_ + +#include // std::string +#include // std::vector + +namespace o2::analysis +{ +namespace hf_cuts_3prongs_alice3 +{ +static constexpr int NBinsPt = 10; +static constexpr int NCutVars = 10; +// default values for the pT bin edges (can be used to configure histogram axis) +// offset by 1 from the bin numbers in cuts array +constexpr double BinsPt[NBinsPt + 1] = { + 0., + 1., + 2., + 3., + 4., + 5., + 6., + 8., + 12., + 24., + 36.}; +const auto vecBinsPt = std::vector{BinsPt, BinsPt + NBinsPt + 1}; + +// default values for the cuts m, ptP, ptK, ptPi, chi2PCA, cosp, dL, dLXY, NdLXY, ImpParXY +constexpr double Cuts[NBinsPt][NCutVars] = {{0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}, /* 0 < pT < 1 */ + {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}, /* 1 < pT < 2 */ + {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}, /* 2 < pT < 3 */ + {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}, /* 3 < pT < 4 */ + {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}, /* 4 < pT < 5 */ + {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}, /* 5 < pT < 6 */ + {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}, /* 6 < pT < 8 */ + {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}, /* 8 < pT < 12 */ + {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}, /* 12 < pT < 24 */ + {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10}}; /* 24 < pT < 36 */ + +// row labels +static const std::vector labelsPt = { + "pT bin 0", + "pT bin 1", + "pT bin 2", + "pT bin 3", + "pT bin 4", + "pT bin 5", + "pT bin 6", + "pT bin 7", + "pT bin 8", + "pT bin 9"}; + +// column labels +static const std::vector labelsCutVar = {"m", "pT prong 0", "pT prong 1", "pT prong 2", "Chi2PCA", "cos pointing angle", "decay length", "decLengthXY", "normDecLXY", "impParXY"}; +} // namespace hf_cuts_3prongs_alice3 + +} // namespace o2::analysis + +#endif // ALICE3_UTILS_UTILSSELECTIONSALICE3_H_