From 04fa7ddf4dfd0b7e53a0d252556f46672e774b2e Mon Sep 17 00:00:00 2001 From: Daiki Sekihata Date: Wed, 20 May 2026 19:02:29 +0200 Subject: [PATCH 1/5] [PWGEM/Dilepton]: update taggingHFE --- PWGEM/Dilepton/DataModel/dileptonTables.h | 16 + PWGEM/Dilepton/TableProducer/CMakeLists.txt | 5 + .../skimmerPrimaryElectronSCT.cxx | 283 +++ PWGEM/Dilepton/Tasks/taggingHFE.cxx | 2 +- PWGEM/Dilepton/Utils/ElectronModule.h | 2240 +++++++++++++++++ PWGEM/Dilepton/Utils/MlResponsePID.h | 147 ++ PWGEM/Dilepton/Utils/MlResponseSCT.h | 264 ++ 7 files changed, 2956 insertions(+), 1 deletion(-) create mode 100644 PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronSCT.cxx create mode 100644 PWGEM/Dilepton/Utils/ElectronModule.h create mode 100644 PWGEM/Dilepton/Utils/MlResponsePID.h create mode 100644 PWGEM/Dilepton/Utils/MlResponseSCT.h diff --git a/PWGEM/Dilepton/DataModel/dileptonTables.h b/PWGEM/Dilepton/DataModel/dileptonTables.h index 9b47e2d4dd4..7bbce9eeef0 100644 --- a/PWGEM/Dilepton/DataModel/dileptonTables.h +++ b/PWGEM/Dilepton/DataModel/dileptonTables.h @@ -691,6 +691,15 @@ DECLARE_SOA_COLUMN(PrefilterBit, pfb, uint8_t); //! DECLARE_SOA_COLUMN(PrefilterBitDerived, pfbderived, uint16_t); //! DECLARE_SOA_COLUMN(ProbElBDT, probElBDT, float); //! +DECLARE_SOA_COLUMN(BDTScorePrompt, bdtScorePrompt, std::vector); //! +DECLARE_SOA_COLUMN(BDTScorePromptHc, bdtScorePromptHc, std::vector); //! +DECLARE_SOA_COLUMN(BDTScoreNonpromptHc, bdtScoreNonpromptHc, std::vector); //! +DECLARE_SOA_COLUMN(BDTScoreHb, bdtScoreHb, std::vector); //! +DECLARE_SOA_COLUMN(HadronType, hadronType, std::vector); //! 0:track, 1:K0S, 2:Lambda, 3:AntiLambda, 4:XiMinus, 5:XiPlus, 6:OmegaMinus, 7:OmegaPlus + +DECLARE_SOA_DYNAMIC_COLUMN(ProbaSCT, probaSCT, [](std::vector p0, std::vector p1, std::vector p2, std::vector p3, std::vector type, int index) -> std::array { return std::array{p0[index], p1[index], p2[index], p3[index], static_cast(type[index])}; }); +DECLARE_SOA_DYNAMIC_COLUMN(NSV, nSV, [](std::vector type) -> size_t { return type.size(); }); + DECLARE_SOA_COLUMN(ITSNSigmaEl, itsNSigmaEl, float); //! DECLARE_SOA_COLUMN(ITSNSigmaMu, itsNSigmaMu, float); //! DECLARE_SOA_COLUMN(ITSNSigmaPi, itsNSigmaPi, float); //! @@ -1007,6 +1016,13 @@ DECLARE_SOA_TABLE(EMPrimaryElectronsPrefilterBitDerived, "AOD", "PRMELPFBDERIVED // iterators using EMPrimaryElectronPrefilterBitDerived = EMPrimaryElectronsPrefilterBitDerived::iterator; +DECLARE_SOA_TABLE(EMPrimaryElectronsBDTSCT, "AOD", "ELBDTSCT", // To be joined with EMPrimaryElectrons table at analysis level. + emprimaryelectron::BDTScorePrompt, emprimaryelectron::BDTScorePromptHc, emprimaryelectron::BDTScoreNonpromptHc, emprimaryelectron::BDTScoreHb, emprimaryelectron::HadronType, + emprimaryelectron::NSV, + emprimaryelectron::ProbaSCT); +// iterators +using EMPrimaryElectronBDTSCT = EMPrimaryElectronsBDTSCT::iterator; + namespace emprimarymuon { DECLARE_SOA_INDEX_COLUMN(EMEvent, emevent); //! diff --git a/PWGEM/Dilepton/TableProducer/CMakeLists.txt b/PWGEM/Dilepton/TableProducer/CMakeLists.txt index 617d07a86d1..0070a556748 100644 --- a/PWGEM/Dilepton/TableProducer/CMakeLists.txt +++ b/PWGEM/Dilepton/TableProducer/CMakeLists.txt @@ -31,6 +31,11 @@ o2physics_add_dpl_workflow(skimmer-primary-electron PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::AnalysisCore O2Physics::MLCore COMPONENT_NAME Analysis) +o2physics_add_dpl_workflow(skimmer-primary-electron-sct + SOURCES skimmerPrimaryElectronSCT.cxx + PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::AnalysisCore O2Physics::MLCore + COMPONENT_NAME Analysis) + o2physics_add_dpl_workflow(skimmer-primary-electron-qc SOURCES skimmerPrimaryElectronQC.cxx PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::AnalysisCore O2Physics::MLCore diff --git a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronSCT.cxx b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronSCT.cxx new file mode 100644 index 00000000000..1df2cf31f2a --- /dev/null +++ b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronSCT.cxx @@ -0,0 +1,283 @@ +// 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. + +/// \brief write relevant information about primary electrons. +/// \author daiki.sekihata@cern.ch + +#include "PWGEM/Dilepton/DataModel/dileptonTables.h" +#include "PWGEM/Dilepton/Utils/ElectronModule.h" +#include "PWGEM/Dilepton/Utils/MlResponseO2Track.h" +#include "PWGEM/Dilepton/Utils/PairUtilities.h" +#include "PWGLF/DataModel/LFStrangenessTables.h" + +#include "Common/Core/TableHelper.h" +#include "Common/Core/trackUtilities.h" +#include "Common/DataModel/CollisionAssociationTables.h" +#include "Common/DataModel/EventSelection.h" +#include "Common/DataModel/PIDResponseTOF.h" +#include "Common/DataModel/PIDResponseTPC.h" +#include "Tools/ML/MlResponse.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include // IWYU pragma: keep (do not replace with Math/Vector4Dfwd.h) +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +// using namespace o2; +// using namespace o2::soa; +// using namespace o2::framework; +// using namespace o2::framework::expressions; +// using namespace o2::constants::physics; +// using namespace o2::common::core; + +struct skimmerPrimaryElectronSCT { + + using MyBCs = o2::soa::Join; + using MyCollisions = o2::soa::Join; + using MyCollisionsWithSWT = o2::soa::Join; + + using MyTracks = o2::soa::Join; + using MyTrack = MyTracks::iterator; + using MyTracksMC = o2::soa::Join; + using MyTrackMC = MyTracksMC::iterator; + + using MyV0s = o2::soa::Join; + using MyCascades = o2::soa::Join; + + struct : o2::framework::ConfigurableGroup { + o2::framework::Configurable ccdburl{"ccdb-url", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; + o2::framework::Configurable grpPath{"grpPath", "GLO/GRP/GRP", "Path of the grp file"}; + o2::framework::Configurable grpmagPath{"grpmagPath", "GLO/Config/GRPMagField", "CCDB path of the GRPMagField object"}; + o2::framework::Configurable lutPath{"lutPath", "GLO/Param/MatLUT", "Path of the Lut parametrization"}; + o2::framework::Configurable mVtxPath{"mVtxPath", "GLO/Calib/MeanVertex", "Path of the mean vertex file"}; + } ccdbConfig; + + o2::framework::Configurable doSCTwithTracks{"doSCTwithTracks", true, "flag to tag electrons with tracks"}; + o2::framework::Configurable doSCTwithV0s{"doSCTwithV0s", false, "flag to tag electrons with v0s"}; + o2::framework::Configurable doSCTwithCascades{"doSCTwithCascades", false, "flag to tag electrons with cascades"}; + + o2::framework::Service ccdb; + o2::ccdb::CcdbApi ccdbApi; + o2::framework::Service mTOFResponse; + + o2::framework::SliceCache cache; + o2::framework::Preslice perCol_track = o2::aod::track::collisionId; + o2::framework::Preslice trackIndicesPerCollision = o2::aod::track_association::collisionId; + o2::framework::Preslice perCol_v0 = o2::aod::v0data::collisionId; + o2::framework::Preslice perCol_casc = o2::aod::cascdata::collisionId; + + void init(o2::framework::InitContext& initContext) + { + mRunNumber = 0; + d_bz = 0; + + LOGF(info, "initializing CCDB"); + ccdb->setURL(ccdbConfig.ccdburl.value); + ccdb->setCaching(true); + ccdb->setLocalObjectValidityChecking(); + ccdb->setFatalWhenNull(false); + // ccdbApi.init(ccdbConfig.ccdburl); + // mTOFResponse->initSetup(ccdb, initContext); + + LOGF(info, "initializing electronModule"); + electronModule.init(cfgEelectronCut, cfgEelectronPFCut, cfgHadronCut, cfgV0Cut, cfgCascadeCut, cfgDFeT, cfgDFeV0, cfgDFeC, initContext, ccdb, mTOFResponse, ccdbConfig.ccdburl.value); + // electronModule.setTOFResponse(mTOFResponse); + electronModule.addHistograms(mRegistry); + // electronModule.doPFB(doPF.value); + + electronModule.doSCTwithTracks(doSCTwithTracks.value); + electronModule.doSCTwithV0s(doSCTwithV0s.value); + electronModule.doSCTwithCascades(doSCTwithCascades.value); + } + + o2::pwgem::dilepton::utils::ElectronModule electronModule; + o2::pwgem::dilepton::utils::ElectronProducts products; + o2::pwgem::dilepton::utils::electronCut cfgEelectronCut; + o2::pwgem::dilepton::utils::electronPFCut cfgEelectronPFCut; + o2::pwgem::dilepton::utils::hadronCut cfgHadronCut; + o2::pwgem::dilepton::utils::v0Cut cfgV0Cut; + o2::pwgem::dilepton::utils::cascadeCut cfgCascadeCut; + o2::pwgem::dilepton::utils::cfgDFeT cfgDFeT; + o2::pwgem::dilepton::utils::cfgDFeV0 cfgDFeV0; + o2::pwgem::dilepton::utils::cfgDFeC cfgDFeC; + + o2::framework::HistogramRegistry mRegistry{"output", {}, o2::framework::OutputObjHandlingPolicy::AnalysisObject, false, false}; + // ---------- for data ---------- + + int mRunNumber{0}; + float d_bz{0}; + o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrLUT; + o2::dataformats::VertexBase mVtx; + // const o2::dataformats::MeanVertexObject* mMeanVtx = nullptr; + o2::base::MatLayerCylSet* lut = nullptr; + + template + void initCCDB(TBC const& bc) + { + if (mRunNumber == bc.runNumber()) { + return; + } + + // load matLUT for this timestamp + if (!lut) { + LOG(info) << "Loading material look-up table for timestamp: " << bc.timestamp(); + lut = o2::base::MatLayerCylSet::rectifyPtrFromFile(ccdb->getForTimeStamp(ccdbConfig.lutPath, bc.timestamp())); + } else { + LOG(info) << "Material look-up table already in place. Not reloading."; + } + + auto run3grp_timestamp = bc.timestamp(); + o2::parameters::GRPMagField* grpmag = 0x0; + o2::parameters::GRPObject* grpo = ccdb->getForTimeStamp(ccdbConfig.grpPath, run3grp_timestamp); + if (grpo) { + o2::base::Propagator::initFieldFromGRP(grpo); + o2::base::Propagator::Instance()->setMatLUT(lut); + // mMeanVtx = ccdb->getForTimeStamp(ccdbConfig.mVtxPath, bc.timestamp()); + // Fetch magnetic field from ccdb for current collision + d_bz = grpo->getNominalL3Field(); + LOG(info) << "Retrieved GRP for timestamp " << run3grp_timestamp << " with magnetic field of " << d_bz << " kZG"; + } else { + grpmag = ccdb->getForTimeStamp(ccdbConfig.grpmagPath, run3grp_timestamp); + if (!grpmag) { + LOG(fatal) << "Got nullptr from CCDB for path " << ccdbConfig.grpmagPath << " of object GRPMagField and " << ccdbConfig.grpPath << " of object GRPObject for timestamp " << run3grp_timestamp; + } + o2::base::Propagator::initFieldFromGRP(grpmag); + o2::base::Propagator::Instance()->setMatLUT(lut); + // mMeanVtx = ccdb->getForTimeStamp(ccdbConfig.mVtxPath, bc.timestamp()); + + // Fetch magnetic field from ccdb for current collision + d_bz = std::lround(5.f * grpmag->getL3Current() / 30000.f); + LOG(info) << "Retrieved GRP for timestamp " << run3grp_timestamp << " with magnetic field of " << d_bz << " kZG"; + } + + mRunNumber = bc.runNumber(); + } + + //! type of V0. 0: built solely for cascades (does not pass standard V0 cut), 1: standard 2, 3: photon-like with TPC-only use. Regular analysis should always use type 1. + o2::framework::expressions::Filter v0Filter = o2::aod::v0data::v0Type == uint8_t(1) && o2::aod::v0data::v0cosPA > cfgV0Cut.cfg_min_cospa&& o2::aod::v0data::dcaV0daughters cfgV0Cut.cfg_min_dcaxy&& nabs(o2::aod::v0data::dcanegtopv) > cfgV0Cut.cfg_min_dcaxy; + using filteredMyV0s = o2::soa::Filtered; + + o2::framework::expressions::Filter cascadeFilter = nabs(o2::aod::cascdata::dcanegtopv) > cfgCascadeCut.cfg_min_dcaxy_v0leg&& nabs(o2::aod::cascdata::dcanegtopv) > cfgCascadeCut.cfg_min_dcaxy_v0leg&& nabs(o2::aod::cascdata::dcabachtopv) > cfgCascadeCut.cfg_min_dcaxy_bachelor&& o2::aod::cascdata::dcacascdaughters < cfgCascadeCut.cfg_max_dcadau&& o2::aod::cascdata::dcaV0daughters < cfgCascadeCut.cfg_max_dcadau_v0; + using filteredMyCascades = o2::soa::Filtered; + + void processRec_SA(MyCollisions const& collisions, MyBCs const& bcs, MyTracks const& tracks, filteredMyV0s const& v0s, filteredMyCascades const& cascades) + { + initCCDB(bcs.begin()); + electronModule.processWithoutTTCA(bcs, collisions, tracks, v0s, cascades, nullptr, nullptr, products, mRegistry); + } + PROCESS_SWITCH(skimmerPrimaryElectronSCT, processRec_SA, "process reconstructed info only", true); // standalone + + void processRec_TTCA(MyCollisions const& collisions, MyBCs const& bcs, MyTracks const& tracks, o2::aod::TrackAssoc const& trackIndices, filteredMyV0s const& v0s, filteredMyCascades const& cascades) + { + initCCDB(bcs.begin()); + electronModule.processWithTTCA(bcs, collisions, tracks, v0s, cascades, trackIndices, nullptr, nullptr, products, mRegistry, cache, perCol_track, trackIndicesPerCollision, perCol_v0, perCol_casc); + } + PROCESS_SWITCH(skimmerPrimaryElectronSCT, processRec_TTCA, "process reconstructed info only", false); // with TTCA + + void processRec_SA_SWT(MyCollisionsWithSWT const& collisions, MyBCs const& bcs, MyTracks const& tracks, filteredMyV0s const& v0s, filteredMyCascades const& cascades) + { + initCCDB(bcs.begin()); + electronModule.processWithoutTTCA(bcs, collisions, tracks, v0s, cascades, nullptr, nullptr, products, mRegistry); + } + PROCESS_SWITCH(skimmerPrimaryElectronSCT, processRec_SA_SWT, "process reconstructed info only", false); // standalone with swt + + void processRec_TTCA_SWT(MyCollisionsWithSWT const& collisions, MyBCs const& bcs, MyTracks const& tracks, o2::aod::TrackAssoc const& trackIndices, filteredMyV0s const& v0s, filteredMyCascades const& cascades) + { + initCCDB(bcs.begin()); + electronModule.processWithTTCA(bcs, collisions, tracks, v0s, cascades, trackIndices, nullptr, nullptr, products, mRegistry, cache, perCol_track, trackIndicesPerCollision, perCol_v0, perCol_casc); + } + PROCESS_SWITCH(skimmerPrimaryElectronSCT, processRec_TTCA_SWT, "process reconstructed info only", false); // with TTCA with swt + + // ---------- for MC ---------- + + void processMC_SA(o2::soa::Join const& collisions, MyBCs const& bcs, MyTracksMC const& tracks, filteredMyV0s const& v0s, filteredMyCascades const& cascades, o2::aod::McCollisions const& mcCollisions, o2::aod::McParticles const& mcParticles) + { + initCCDB(bcs.begin()); + electronModule.processWithoutTTCA(bcs, collisions, tracks, v0s, cascades, mcCollisions, mcParticles, products, mRegistry); + } + PROCESS_SWITCH(skimmerPrimaryElectronSCT, processMC_SA, "process reconstructed and MC info ", false); // without TTCA in MC + + void processMC_TTCA(o2::soa::Join const& collisions, MyBCs const& bcs, MyTracksMC const& tracks, o2::aod::TrackAssoc const& trackIndices, filteredMyV0s const& v0s, filteredMyCascades const& cascades, o2::aod::McCollisions const& mcCollisions, o2::aod::McParticles const& mcParticles) + { + initCCDB(bcs.begin()); + electronModule.processWithTTCA(bcs, collisions, tracks, v0s, cascades, trackIndices, mcCollisions, mcParticles, products, mRegistry, cache, perCol_track, trackIndicesPerCollision, perCol_v0, perCol_casc); + } + PROCESS_SWITCH(skimmerPrimaryElectronSCT, processMC_TTCA, "process reconstructed info only", false); // with TTCA in MC +}; + +struct associateAmbiguousElectron { + o2::framework::Produces em_amb_ele_ids; + + o2::framework::SliceCache cache; + o2::framework::PresliceUnsorted perTrack = o2::aod::emprimaryelectron::trackId; + std::vector ambele_self_Ids; + + void process(o2::aod::EMPrimaryElectrons const& electrons) + { + for (const auto& electron : electrons) { + auto electrons_with_same_trackId = electrons.sliceBy(perTrack, electron.trackId()); + ambele_self_Ids.reserve(electrons_with_same_trackId.size()); + for (const auto& amb_ele : electrons_with_same_trackId) { + if (amb_ele.globalIndex() == electron.globalIndex()) { // don't store myself. + continue; + } + ambele_self_Ids.emplace_back(amb_ele.globalIndex()); + } + em_amb_ele_ids(ambele_self_Ids); + ambele_self_Ids.clear(); + ambele_self_Ids.shrink_to_fit(); + } + } +}; +o2::framework::WorkflowSpec defineDataProcessing(o2::framework::ConfigContext const& cfgc) +{ + o2::pid::tof::TOFResponseImpl::metadataInfo.initMetadata(cfgc); + return o2::framework::WorkflowSpec{ + adaptAnalysisTask(cfgc, o2::framework::TaskName{"skimmer-primary-electron-sct"}), + adaptAnalysisTask(cfgc, o2::framework::TaskName{"associate-ambiguous-electron"})}; +} diff --git a/PWGEM/Dilepton/Tasks/taggingHFE.cxx b/PWGEM/Dilepton/Tasks/taggingHFE.cxx index 893767d09ee..163b6b3089e 100644 --- a/PWGEM/Dilepton/Tasks/taggingHFE.cxx +++ b/PWGEM/Dilepton/Tasks/taggingHFE.cxx @@ -440,7 +440,7 @@ struct taggingHFE { fRegistry.add("Generated/D0/hsAcc", "pT-#eta acc.;p_{T,l} (GeV/c);p_{T,K} (GeV/c);#eta_{l};#eta_{K};", kTHnSparseF, {{100, 0, 10}, {100, 0, 10}, {100, -5, +5}, {100, -5, +5}}, false); fRegistry.add("Generated/Lc/hsAcc", "pT-#eta acc.;p_{T,l} (GeV/c);p_{T,#Lambda} (GeV/c);#eta_{l};#eta_{#Lambda};", kTHnSparseF, {{100, 0, 10}, {100, 0, 10}, {100, -5, +5}, {100, -5, +5}}, false); - fRegistry.add("Electron/hs", "hs;p_{T} (GeV/c);#eta;#varphi (rad.)", kTHnSparseF, {{100, 0, 10}, {40, -1, 1}, {36, 0, 2 * M_PI}}, false); + fRegistry.add("Electron/hs", "hs;p_{T} (GeV/c);#eta;#varphi (rad.)", kTHnSparseF, {{100, 0, 10}, {80, -2, 2}, {36, 0, 2 * M_PI}}, false); fRegistry.add("Electron/hDCA", "DCA xy vs. z;DCA_{xy} (cm);DCA_{z} (cm)", kTH2F, {{200, -1, 1}, {200, -1, 1}}, false); fRegistry.add("Electron/hTPCdEdx", "TPC dE/dx vs. pin;p_{in} (GeV/c);TPC dE/dx", kTH2F, {{1000, 0, 10}, {200, 0, 200}}, false); fRegistry.add("Electron/hTOFbeta", "TOF #beta vs. p;p_{pv} (GeV/c);TOF #beta", kTH2F, {{1000, 0, 10}, {600, 0, 1.2}}, false); diff --git a/PWGEM/Dilepton/Utils/ElectronModule.h b/PWGEM/Dilepton/Utils/ElectronModule.h new file mode 100644 index 00000000000..9abd547c51e --- /dev/null +++ b/PWGEM/Dilepton/Utils/ElectronModule.h @@ -0,0 +1,2240 @@ +// 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. + +/// \brief write relevant information about primary electrons. +/// \author daiki.sekihata@cern.ch + +#ifndef PWGEM_DILEPTON_UTILS_ELECTRONMODULE_H_ +#define PWGEM_DILEPTON_UTILS_ELECTRONMODULE_H_ + +#include "PWGEM/Dilepton/DataModel/dileptonTables.h" +#include "PWGEM/Dilepton/Utils/MlResponsePID.h" +#include "PWGEM/Dilepton/Utils/MlResponseSCT.h" +#include "PWGEM/Dilepton/Utils/PairUtilities.h" +#include "PWGEM/Dilepton/Utils/SemiCharmTag.h" +#include "PWGLF/DataModel/LFStrangenessTables.h" + +#include "Common/Core/RecoDecay.h" +#include "Common/Core/TableHelper.h" +#include "Common/Core/trackUtilities.h" +#include "Common/DataModel/CollisionAssociationTables.h" +#include "Common/DataModel/EventSelection.h" +#include "Common/DataModel/PIDResponseTOF.h" +#include "Common/DataModel/PIDResponseTPC.h" +#include "Tools/ML/MlResponse.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include // IWYU pragma: keep (do not replace with Math/Vector4Dfwd.h) +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace o2 +{ +namespace pwgem::dilepton::utils +{ + +struct ElectronProducts : o2::framework::ProducesGroup { + o2::framework::Produces electronTable; + o2::framework::Produces electronCovTable; + o2::framework::Produces electronPFTable; + o2::framework::Produces sctTable; +}; + +struct electronCut : o2::framework::ConfigurableGroup { + std::string prefix = "electronCut"; + o2::framework::Configurable min_ncluster_tpc{"min_ncluster_tpc", 0, "min ncluster tpc"}; + o2::framework::Configurable mincrossedrows{"mincrossedrows", 70, "min. crossed rows"}; + o2::framework::Configurable min_tpc_cr_findable_ratio{"min_tpc_cr_findable_ratio", 0.8, "min. TPC Ncr/Nf ratio"}; + o2::framework::Configurable min_ncluster_its{"min_ncluster_its", 4, "min ncluster its"}; + o2::framework::Configurable min_ncluster_itsib{"min_ncluster_itsib", 1, "min ncluster itsib"}; + o2::framework::Configurable minchi2tpc{"minchi2tpc", 0.0, "min. chi2/NclsTPC"}; + o2::framework::Configurable minchi2its{"minchi2its", 0.0, "min. chi2/NclsITS"}; + o2::framework::Configurable maxchi2tpc{"maxchi2tpc", 5.0, "max. chi2/NclsTPC"}; + o2::framework::Configurable maxchi2its{"maxchi2its", 6.0, "max. chi2/NclsITS"}; + o2::framework::Configurable minpt{"minpt", 0.05, "min pt"}; + o2::framework::Configurable maxeta{"maxeta", 0.9, "max eta"}; + o2::framework::Configurable dca_xy_max{"dca_xy_max", 1.0, "max DCAxy in cm"}; + o2::framework::Configurable dca_z_max{"dca_z_max", 1.0, "max DCAz in cm"}; + o2::framework::Configurable minTPCNsigmaEl{"minTPCNsigmaEl", -2.5, "min. TPC n sigma for electron inclusion"}; + o2::framework::Configurable maxTPCNsigmaEl{"maxTPCNsigmaEl", 3.5, "max. TPC n sigma for electron inclusion"}; + o2::framework::Configurable maxTOFNsigmaEl{"maxTOFNsigmaEl", 3.5, "max. TOF n sigma for electron inclusion"}; + o2::framework::Configurable maxTPCNsigmaPi{"maxTPCNsigmaPi", 2.5, "max. TPC n sigma for pion exclusion"}; + o2::framework::Configurable minTPCNsigmaPi{"minTPCNsigmaPi", -1e+10, "min. TPC n sigma for pion exclusion"}; // set to -2 for lowB, -1e+10 for nominalB + o2::framework::Configurable maxTPCNsigmaKa{"maxTPCNsigmaKa", 2.5, "max. TPC n sigma for kaon exclusion"}; + o2::framework::Configurable minTPCNsigmaKa{"minTPCNsigmaKa", -2.5, "min. TPC n sigma for kaon exclusion"}; + o2::framework::Configurable maxTPCNsigmaPr{"maxTPCNsigmaPr", 2.5, "max. TPC n sigma for proton exclusion"}; + o2::framework::Configurable minTPCNsigmaPr{"minTPCNsigmaPr", -2.5, "min. TPC n sigma for proton exclusion"}; + o2::framework::Configurable requireTOF{"requireTOF", false, "require TOF hit"}; + o2::framework::Configurable min_pin_for_pion_rejection{"min_pin_for_pion_rejection", 0.0, "pion rejection is applied above this pin"}; // this is used only in TOFreq + o2::framework::Configurable max_pin_for_pion_rejection{"max_pin_for_pion_rejection", 0.5, "pion rejection is applied below this pin"}; + o2::framework::Configurable max_frac_shared_clusters_tpc{"max_frac_shared_clusters_tpc", 999.f, "max fraction of shared clusters in TPC"}; + o2::framework::Configurable maxMeanITSClusterSize{"maxMeanITSClusterSize", 16, "max x cos(lambda)"}; + o2::framework::Configurable storeOnlyTrueElectronMC{"storeOnlyTrueElectronMC", false, "Flag to store only true electron in MC"}; + o2::framework::Configurable minNelectron{"minNelectron", 0, "min number of electron candidates per collision"}; + o2::framework::Configurable includeITSsa{"includeITSsa", false, "Flag to include ITSsa tracks only for MC. switch ON only if needed."}; + o2::framework::Configurable useTOFNSigmaDeltaBC{"useTOFNSigmaDeltaBC", false, "Flag to shift delta BC for TOF n sigma (only with TTCA)"}; + + // configuration for PID ML + o2::framework::Configurable usePIDML{"usePIDML", false, "Flag to use PID ML"}; + o2::framework::Configurable> onnxFileNames{"onnxFileNames", std::vector{"filename"}, "ONNX file names for each bin (if not from CCDB full path)"}; + o2::framework::Configurable> onnxPathsCCDB{"onnxPathsCCDB", std::vector{"path"}, "Paths of models on CCDB"}; + o2::framework::Configurable> binsMl{"binsMl", std::vector{0.1, 0.15, 0.2, 0.25, 0.4, 0.8, 1.6, 2.0, 20}, "Bin limits for ML application"}; + o2::framework::Configurable> cutsMl{"cutsMl", std::vector{0.95, 0.95, 0.7, 0.7, 0.8, 0.8, 0.7, 0.7}, "ML cuts per bin"}; + o2::framework::Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"tpcInnerParam", "tpcNClsFound", "tpcChi2NCl", "tpcNSigmaEl", "tofNSigmaEl", "meanClusterSizeITSobCosTgl"}, "Names of ML model input features"}; + o2::framework::Configurable nameBinningFeature{"nameBinningFeature", "tpcInnerParam", "Names of ML model binning feature"}; + o2::framework::Configurable loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; + o2::framework::Configurable enableOptimizations{"enableOptimizations", false, "Enables the ONNX extended model-optimization: sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED)"}; +}; + +struct electronPFCut : o2::framework::ConfigurableGroup { + std::string prefix = "electronPFCut"; + o2::framework::Configurable min_ncluster_tpc{"min_ncluster_tpc", 0, "min ncluster tpc"}; + o2::framework::Configurable mincrossedrows{"mincrossedrows", 70, "min. crossed rows"}; + o2::framework::Configurable min_tpc_cr_findable_ratio{"min_tpc_cr_findable_ratio", 0.8, "min. TPC Ncr/Nf ratio"}; + o2::framework::Configurable min_ncluster_its{"min_ncluster_its", 4, "min ncluster its"}; + o2::framework::Configurable min_ncluster_itsib{"min_ncluster_itsib", 1, "min ncluster itsib"}; + o2::framework::Configurable minchi2tpc{"minchi2tpc", 0.0, "min. chi2/NclsTPC"}; + o2::framework::Configurable minchi2its{"minchi2its", 0.0, "min. chi2/NclsITS"}; + o2::framework::Configurable maxchi2tpc{"maxchi2tpc", 5.0, "max. chi2/NclsTPC"}; + o2::framework::Configurable maxchi2its{"maxchi2its", 36.0, "max. chi2/NclsITS"}; + o2::framework::Configurable minpt{"minpt", 0.05, "min pt"}; + o2::framework::Configurable maxeta{"maxeta", 0.9, "max eta"}; + o2::framework::Configurable dca_xy_max{"dca_xy_max", 1.0, "max DCAxy in cm"}; + o2::framework::Configurable dca_z_max{"dca_z_max", 1.0, "max DCAz in cm"}; + o2::framework::Configurable minTPCNsigmaEl{"minTPCNsigmaEl", -2.0, "min. TPC n sigma for electron inclusion"}; + o2::framework::Configurable maxTPCNsigmaEl{"maxTPCNsigmaEl", 3.0, "max. TPC n sigma for electron inclusion"}; + o2::framework::Configurable maxTOFNsigmaEl{"maxTOFNsigmaEl", 1e+10, "max. TOF n sigma for electron inclusion"}; + o2::framework::Configurable maxTPCNsigmaPi{"maxTPCNsigmaPi", 0.0, "max. TPC n sigma for pion exclusion"}; + o2::framework::Configurable minTPCNsigmaPi{"minTPCNsigmaPi", 0.0, "min. TPC n sigma for pion exclusion"}; // set to -2 for lowB, -1e+10 for nominalB + o2::framework::Configurable maxTPCNsigmaKa{"maxTPCNsigmaKa", 0.0, "max. TPC n sigma for kaon exclusion"}; + o2::framework::Configurable minTPCNsigmaKa{"minTPCNsigmaKa", 0.0, "min. TPC n sigma for kaon exclusion"}; + o2::framework::Configurable maxTPCNsigmaPr{"maxTPCNsigmaPr", 0.0, "max. TPC n sigma for proton exclusion"}; + o2::framework::Configurable minTPCNsigmaPr{"minTPCNsigmaPr", 0.0, "min. TPC n sigma for proton exclusion"}; + // o2::framework::Configurable requireTOF{"requireTOF", false, "require TOF hit"}; + o2::framework::Configurable min_pin_for_pion_rejection{"min_pin_for_pion_rejection", 0.0, "pion rejection is applied above this pin"}; // this is used only in TOFreq + o2::framework::Configurable max_pin_for_pion_rejection{"max_pin_for_pion_rejection", 0.5, "pion rejection is applied below this pin"}; + o2::framework::Configurable max_frac_shared_clusters_tpc{"max_frac_shared_clusters_tpc", 999.f, "max fraction of shared clusters in TPC"}; + o2::framework::Configurable maxMeanITSClusterSize{"maxMeanITSClusterSize", 16, "max x cos(lambda)"}; + + // configuration for PID ML + o2::framework::Configurable> binsMl{"binsMl", std::vector{0.1, 0.15, 0.2, 0.25, 0.4, 0.8, 1.6, 2.0, 20}, "Bin limits for ML application"}; + o2::framework::Configurable> cutsMl{"cutsMl", std::vector{0.9, 0.9, 0.65, 0.65, 0.75, 0.75, 0.65, 0.65}, "ML cuts per bin"}; + + // for pair + o2::framework::Configurable slope{"slope", 0.0185, "slope for m vs. phiv"}; + o2::framework::Configurable intercept{"intercept", -0.0280, "intercept for m vs. phiv"}; + o2::framework::Configurable doPF{"doPF", false, "flag to set pion prefilter"}; +}; + +struct hadronCut : o2::framework::ConfigurableGroup { + std::string prefix = "hadronCut"; + o2::framework::Configurable cfg_min_pt_track{"cfg_min_pt_track", 0.01, "min pT for single track"}; + o2::framework::Configurable cfg_max_pt_track{"cfg_max_pt_track", 1e+10, "max pT for single track"}; + o2::framework::Configurable cfg_min_eta_track{"cfg_min_eta_track", -0.9, "min eta for single track"}; + o2::framework::Configurable cfg_max_eta_track{"cfg_max_eta_track", +0.9, "max eta for single track"}; + o2::framework::Configurable cfg_min_cr2findable_ratio_tpc{"cfg_min_cr2findable_ratio_tpc", 0.f, "min. TPC Ncr/Nf ratio"}; + o2::framework::Configurable cfg_max_frac_shared_clusters_tpc{"cfg_max_frac_shared_clusters_tpc", 999.f, "max fraction of shared clusters in TPC"}; + o2::framework::Configurable cfg_min_ncrossedrows_tpc{"cfg_min_ncrossedrows_tpc", 70, "min ncrossed rows"}; + o2::framework::Configurable cfg_min_ncluster_tpc{"cfg_min_ncluster_tpc", 0, "min ncluster tpc"}; + o2::framework::Configurable cfg_min_ncluster_its{"cfg_min_ncluster_its", 4, "min ncluster its"}; + o2::framework::Configurable cfg_min_ncluster_itsib{"cfg_min_ncluster_itsib", 1, "min ncluster itsib"}; + o2::framework::Configurable cfg_max_chi2tpc{"cfg_max_chi2tpc", 4.0, "max chi2/NclsTPC"}; + o2::framework::Configurable cfg_max_chi2its{"cfg_max_chi2its", 36.0, "max chi2/NclsITS"}; + o2::framework::Configurable cfg_max_dcaxy{"cfg_max_dcaxy", 1.0, "max dca XY for single track in cm"}; + o2::framework::Configurable cfg_max_dcaz{"cfg_max_dcaz", 1.0, "max dca Z for single track in cm"}; + o2::framework::Configurable cfg_min_TPCNsigmaPi{"cfg_min_TPCNsigmaPi", -3, "min n sigma pi in TPC"}; + o2::framework::Configurable cfg_max_TPCNsigmaPi{"cfg_max_TPCNsigmaPi", +3, "max n sigma pi in TPC"}; + o2::framework::Configurable cfg_min_TOFNsigmaPi{"cfg_min_TOFNsigmaPi", -3, "min n sigma pi in TOF"}; + o2::framework::Configurable cfg_max_TOFNsigmaPi{"cfg_max_TOFNsigmaPi", +3, "max n sigma pi in TOF"}; + o2::framework::Configurable cfg_min_TPCNsigmaKa{"cfg_min_TPCNsigmaKa", -3, "min n sigma ka in TPC"}; + o2::framework::Configurable cfg_max_TPCNsigmaKa{"cfg_max_TPCNsigmaKa", +3, "max n sigma ka in TPC"}; + o2::framework::Configurable cfg_min_TOFNsigmaKa{"cfg_min_TOFNsigmaKa", -3, "min n sigma ka in TOF"}; + o2::framework::Configurable cfg_max_TOFNsigmaKa{"cfg_max_TOFNsigmaKa", +3, "max n sigma ka in TOF"}; + o2::framework::Configurable cfg_min_TPCNsigmaPr{"cfg_min_TPCNsigmaPr", -3, "min n sigma pr in TPC"}; + o2::framework::Configurable cfg_max_TPCNsigmaPr{"cfg_max_TPCNsigmaPr", +3, "max n sigma pr in TPC"}; + o2::framework::Configurable cfg_min_TOFNsigmaPr{"cfg_min_TOFNsigmaPr", -3, "min n sigma pr in TOF"}; + o2::framework::Configurable cfg_max_TOFNsigmaPr{"cfg_max_TOFNsigmaPr", +3, "max n sigma pr in TOF"}; + o2::framework::Configurable requirePiKaPr{"requirePiKaPr", true, "require hadron to be pion or kaon or proton"}; +}; + +struct v0Cut : o2::framework::ConfigurableGroup { + std::string prefix = "v0Cut"; + o2::framework::Configurable cfg_min_mass_k0s{"cfg_min_mass_k0s", 0.48, "min mass for K0S"}; + o2::framework::Configurable cfg_max_mass_k0s{"cfg_max_mass_k0s", 0.51, "max mass for K0S"}; + o2::framework::Configurable cfg_min_mass_k0s_veto{"cfg_min_mass_k0s_veto", 0.48, "min mass for K0S veto for Lambda"}; + o2::framework::Configurable cfg_max_mass_k0s_veto{"cfg_max_mass_k0s_veto", 0.51, "max mass for K0S veto for Lambda"}; + o2::framework::Configurable cfg_min_mass_lambda{"cfg_min_mass_lambda", 1.11, "min mass for Lambda"}; + o2::framework::Configurable cfg_max_mass_lambda{"cfg_max_mass_lambda", 1.12, "max mass for Lambda"}; + o2::framework::Configurable cfg_min_mass_lambda_veto{"cfg_min_mass_lambda_veto", 1.11, "min mass for Lambda veto for K0S"}; + o2::framework::Configurable cfg_max_mass_lambda_veto{"cfg_max_mass_lambda_veto", 1.12, "max mass for Lambda veto for K0S"}; + o2::framework::Configurable cfg_min_cospa{"cfg_min_cospa", 0.95, "min cospa for v0"}; + o2::framework::Configurable cfg_max_dca2legs{"cfg_max_dca2legs", 0.1, "max distance between 2 legs for v0"}; + o2::framework::Configurable cfg_min_radius{"cfg_min_radius", 0.1, "min rxy for v"}; + o2::framework::Configurable cfg_min_cr2findable_ratio_tpc{"cfg_min_cr2findable_ratio_tpc", 0.f, "min. TPC Ncr/Nf ratio"}; + o2::framework::Configurable cfg_max_frac_shared_clusters_tpc{"cfg_max_frac_shared_clusters_tpc", 999.f, "max fraction of shared clusters in TPC"}; + o2::framework::Configurable cfg_min_ncrossedrows_tpc{"cfg_min_ncrossedrows_tpc", 70, "min ncrossed rows"}; + o2::framework::Configurable cfg_min_ncluster_tpc{"cfg_min_ncluster_tpc", 0, "min ncluster tpc"}; + o2::framework::Configurable cfg_max_chi2tpc{"cfg_max_chi2tpc", 4.0, "max chi2/NclsTPC"}; + o2::framework::Configurable cfg_max_chi2its{"cfg_max_chi2its", 36.0, "max chi2/NclsITS"}; + o2::framework::Configurable cfg_min_ncluster_its{"cfg_min_ncluster_its", 2, "min ncluster its"}; + o2::framework::Configurable cfg_min_ncluster_itsib{"cfg_min_ncluster_itsib", 0, "min ncluster itsib"}; + o2::framework::Configurable cfg_min_dcaxy{"cfg_min_dcaxy", 0.1, "min dca XY for v0 legs in cm"}; + + o2::framework::Configurable cfg_max_alpha_veto{"cfg_max_alpha_veto", 0.95, "max alpha for photon conversion rejection"}; + o2::framework::Configurable cfg_max_qt_veto{"cfg_max_qt_veto", 0.01, "max qT for photon conversion rejection"}; + + // for both v0 and cascade + o2::framework::Configurable cfg_min_TPCNsigmaPi{"cfg_min_TPCNsigmaPi", -3, "min n sigma pi in TPC"}; + o2::framework::Configurable cfg_max_TPCNsigmaPi{"cfg_max_TPCNsigmaPi", +3, "max n sigma pi in TPC"}; + o2::framework::Configurable cfg_min_TPCNsigmaKa{"cfg_min_TPCNsigmaKa", -3, "min n sigma ka in TPC"}; + o2::framework::Configurable cfg_max_TPCNsigmaKa{"cfg_max_TPCNsigmaKa", +3, "max n sigma ka in TPC"}; + o2::framework::Configurable cfg_min_TPCNsigmaPr{"cfg_min_TPCNsigmaPr", -3, "min n sigma pr in TPC"}; + o2::framework::Configurable cfg_max_TPCNsigmaPr{"cfg_max_TPCNsigmaPr", +3, "max n sigma pr in TPC"}; + // o2::framework::Configurable cfg_min_TOFNsigmaPi{"cfg_min_TOFNsigmaPi", -3, "min n sigma pi in TOF"}; + // o2::framework::Configurable cfg_max_TOFNsigmaPi{"cfg_max_TOFNsigmaPi", +3, "max n sigma pi in TOF"}; + // o2::framework::Configurable cfg_min_TOFNsigmaKa{"cfg_min_TOFNsigmaKa", -3, "min n sigma ka in TOF"}; + // o2::framework::Configurable cfg_max_TOFNsigmaKa{"cfg_max_TOFNsigmaKa", +3, "max n sigma ka in TOF"}; + // o2::framework::Configurable cfg_min_TOFNsigmaPr{"cfg_min_TOFNsigmaPr", -3, "min n sigma pr in TOF"}; + // o2::framework::Configurable cfg_max_TOFNsigmaPr{"cfg_max_TOFNsigmaPr", +3, "max n sigma pr in TOF"}; + // o2::framework::Configurable applyTOFif{"applyTOFif", false, "apply TOFif for hadron identification"}; +}; + +struct cascadeCut : o2::framework::ConfigurableGroup { + std::string prefix = "cascadeCut"; + o2::framework::Configurable cfg_min_mass_lambda{"cfg_min_mass_lambda", 1.11, "min mass for lambda in cascade"}; + o2::framework::Configurable cfg_max_mass_lambda{"cfg_max_mass_lambda", 1.12, "max mass for lambda in cascade"}; + o2::framework::Configurable cfg_min_mass_Xi{"cfg_min_mass_Xi", 1.314, "min mass for Xi"}; + o2::framework::Configurable cfg_max_mass_Xi{"cfg_max_mass_Xi", 1.328, "max mass for Xi"}; + o2::framework::Configurable cfg_min_mass_Xi_veto{"cfg_min_mass_Xi_veto", 1.31, "min mass for Xi veto"}; + o2::framework::Configurable cfg_max_mass_Xi_veto{"cfg_max_mass_Xi_veto", 1.33, "max mass for Xi veto"}; + o2::framework::Configurable cfg_min_mass_Omega{"cfg_min_mass_Omega", 1.668, "min mass for Omega"}; + o2::framework::Configurable cfg_max_mass_Omega{"cfg_max_mass_Omega", 1.678, "max mass for Omega"}; + o2::framework::Configurable cfg_min_mass_Omega_veto{"cfg_min_mass_Omega_veto", 1.665, "min mass for Omega veto"}; + o2::framework::Configurable cfg_max_mass_Omega_veto{"cfg_max_mass_Omega_veto", 1.680, "max mass for Omega veto"}; + o2::framework::Configurable cfg_min_cospa_v0{"cfg_min_cospa_v0", 0.95, "minimum V0 CosPA in cascade"}; + o2::framework::Configurable cfg_max_dcadau_v0{"cfg_max_dcadau_v0", 0.1, "max distance between V0 Daughters in cascade"}; + o2::framework::Configurable cfg_min_cospa{"cfg_min_cospa", 0.95, "minimum cascade CosPA"}; + o2::framework::Configurable cfg_max_dcadau{"cfg_max_dcadau", 0.1, "max distance between bachelor and V0"}; + o2::framework::Configurable cfg_min_rxy_v0{"cfg_min_rxy_v0", 0.1, "minimum V0 rxy in cascade"}; + o2::framework::Configurable cfg_min_rxy{"cfg_min_rxy", 0.1, "minimum V0 rxy in cascade"}; + o2::framework::Configurable cfg_min_dcaxy_v0leg{"cfg_min_dcaxy_v0leg", 0.1, "min dca XY for v0 legs in cm"}; + o2::framework::Configurable cfg_min_dcaxy_bachelor{"cfg_min_dcaxy_bachelor", 0.05, "min dca XY for bachelor in cm"}; + o2::framework::Configurable cfg_min_dcaxy_v0{"cfg_min_dcaxy_v0", 0.0, "min dca XY for V0 in cm"}; +}; + +struct cfgDFeT : o2::framework::ConfigurableGroup { + std::string prefix = "cfgDFeT"; + o2::framework::Configurable useAbsDCA{"useAbsDCA", true, "Minimise abs. distance rather than chi2"}; + o2::framework::Configurable useWeightedFinalPCA{"useWeightedFinalPCA", false, "Recalculate vertex position using track covariances, effective only if useAbsDCA is true"}; + + // configuration for PID ML + o2::framework::Configurable useML{"useML", false, "Flag to use PID ML"}; + o2::framework::Configurable> onnxFileNames{"onnxFileNames", std::vector{"filename"}, "ONNX file names for each bin (if not from CCDB full path)"}; + o2::framework::Configurable> onnxPathsCCDB{"onnxPathsCCDB", std::vector{"path"}, "Paths of models on CCDB"}; + o2::framework::Configurable> binsMl{"binsMl", std::vector{0.1, 0.2, 0.4, 0.8, 1.0, 2.0, 4, 20}, "Bin limits for ML application"}; + // o2::framework::Configurable> cutsMl{"cutsMl", std::vector{0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9}, "ML cuts per bin"}; + o2::framework::Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"ptH", "impPar3DHinSigma", "tpcNSigmaKa", "signedMassLH", "cpa", "cpaXY", "dcaLH", "impPar3DinSigma", "decayLength3DinSigma", "decayLengthXYinSigma"}, "Names of ML model input features"}; + o2::framework::Configurable nameBinningFeature{"nameBinningFeature", "ptL", "Names of ML model binning feature"}; + o2::framework::Configurable loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; + o2::framework::Configurable enableOptimizations{"enableOptimizations", false, "Enables the ONNX extended model-optimization: sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED)"}; +}; + +struct cfgDFeV0 : o2::framework::ConfigurableGroup { + std::string prefix = "cfgDFeV0"; + o2::framework::Configurable useAbsDCA{"useAbsDCA", true, "Minimise abs. distance rather than chi2"}; + o2::framework::Configurable useWeightedFinalPCA{"useWeightedFinalPCA", false, "Recalculate vertex position using track covariances, effective only if useAbsDCA is true"}; + // configuration for PID ML + o2::framework::Configurable useML{"useML", false, "Flag to use PID ML"}; + o2::framework::Configurable> onnxFileNames{"onnxFileNames", std::vector{"filename"}, "ONNX file names for each bin (if not from CCDB full path)"}; + o2::framework::Configurable> onnxPathsCCDB{"onnxPathsCCDB", std::vector{"path"}, "Paths of models on CCDB"}; + o2::framework::Configurable> binsMl{"binsMl", std::vector{0.1, 0.2, 0.4, 0.8, 1.0, 2.0, 4, 20}, "Bin limits for ML application"}; + // o2::framework::Configurable> cutsMl{"cutsMl", std::vector{0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9}, "ML cuts per bin"}; + o2::framework::Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"ptH", "impPar3DHinSigma", "massLH", "cpa", "cpaXY", "dcaLH", "impPar3DinSigma", "decayLength3DinSigma", "decayLengthXYinSigma"}, "Names of ML model input features"}; + o2::framework::Configurable nameBinningFeature{"nameBinningFeature", "ptL", "Names of ML model binning feature"}; + o2::framework::Configurable loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; + o2::framework::Configurable enableOptimizations{"enableOptimizations", false, "Enables the ONNX extended model-optimization: sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED)"}; +}; + +struct cfgDFeC : o2::framework::ConfigurableGroup { + std::string prefix = "cfgDFeC"; + o2::framework::Configurable useAbsDCA{"useAbsDCA", true, "Minimise abs. distance rather than chi2"}; + o2::framework::Configurable useWeightedFinalPCA{"useWeightedFinalPCA", false, "Recalculate vertex position using track covariances, effective only if useAbsDCA is true"}; + // configuration for PID ML + o2::framework::Configurable useML{"useML", false, "Flag to use PID ML"}; + o2::framework::Configurable> onnxFileNames{"onnxFileNames", std::vector{"filename"}, "ONNX file names for each bin (if not from CCDB full path)"}; + o2::framework::Configurable> onnxPathsCCDB{"onnxPathsCCDB", std::vector{"path"}, "Paths of models on CCDB"}; + o2::framework::Configurable> binsMl{"binsMl", std::vector{0.1, 0.2, 0.4, 0.8, 1.0, 2.0, 4, 20}, "Bin limits for ML application"}; + // o2::framework::Configurable> cutsMl{"cutsMl", std::vector{0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9}, "ML cuts per bin"}; + o2::framework::Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"ptH", "impPar3DHinSigma", "massLH", "cpa", "cpaXY", "dcaLH", "impPar3DinSigma", "decayLength3DinSigma", "decayLengthXYinSigma"}, "Names of ML model input features"}; + o2::framework::Configurable nameBinningFeature{"nameBinningFeature", "ptL", "Names of ML model binning feature"}; + o2::framework::Configurable loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; + o2::framework::Configurable enableOptimizations{"enableOptimizations", false, "Enables the ONNX extended model-optimization: sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED)"}; +}; + +class ElectronModule +{ + public: + ElectronModule() + { + // constructor + } + ~ElectronModule() + { + // destructor + } + + struct looseElectron { + float pt{1e+10}; + float eta{1e+10}; + float phi{1e+10}; + }; + + template + void init(TElectronCut const& eCut, TElectronPFCut const& ePFCut, THadronCut const& hCut, TV0Cut const& v0Cut, TCascadeCut const& cascadeCut, TDFConfigET const& cfgET, TDFConfigEV0 const& cfgEV0, TDFConfigEC const& cfgEC, TInitContext& initContext, TCCDB& ccdb, TTOFResponse const& tofResponse, std::string const& ccdburl) + { + mRunNumber = 0; + d_bz = 0; + + ccdbApi.init(ccdburl); + + LOGF(info, "intializing configurations"); + fElectronCut = eCut; + fElectronPFCut = ePFCut; + fHadronCut = hCut; + fV0Cut = v0Cut; + fCascadeCut = cascadeCut; + fConfigDFeT = cfgET; + fConfigDFeV0 = cfgEV0; + fConfigDFeC = cfgEC; + + LOGF(info, "intializing TOFResponse"); + mTOFResponse = tofResponse; + // mTOFResponse->initSetup(&ccdb->instance(), initContext); + mTOFResponse->initSetup(ccdb, initContext); + + dfeT.setPropagateToPCA(true); + dfeT.setMaxR(200.f); + dfeT.setMinParamChange(1e-3); + dfeT.setMinRelChi2Change(0.9); + dfeT.setMaxDZIni(1e9); + dfeT.setMaxChi2(1e9); + dfeT.setUseAbsDCA(fConfigDFeT.useAbsDCA); + dfeT.setWeightedFinalPCA(fConfigDFeT.useWeightedFinalPCA); + dfeT.setMatCorrType(matCorr); + + dfeV0.setPropagateToPCA(true); + dfeV0.setMaxR(200.f); + dfeV0.setMinParamChange(1e-3); + dfeV0.setMinRelChi2Change(0.9); + dfeV0.setMaxDZIni(1e9); + dfeV0.setMaxChi2(1e9); + dfeV0.setUseAbsDCA(fConfigDFeV0.useAbsDCA); + dfeV0.setWeightedFinalPCA(fConfigDFeV0.useWeightedFinalPCA); + dfeV0.setMatCorrType(matCorr); + + dfeC.setPropagateToPCA(true); + dfeC.setMaxR(200.f); + dfeC.setMinParamChange(1e-3); + dfeC.setMinRelChi2Change(0.9); + dfeC.setMaxDZIni(1e9); + dfeC.setMaxChi2(1e9); + dfeC.setUseAbsDCA(fConfigDFeC.useAbsDCA); + dfeC.setWeightedFinalPCA(fConfigDFeC.useWeightedFinalPCA); + dfeC.setMatCorrType(matCorr); + } + + template + void initCCDB(TBC const& bc) + { + if (mRunNumber == bc.runNumber()) { + return; + } + + d_bz = o2::base::Propagator::Instance()->getNominalBz(); + LOG(info) << "Configuring for timestamp " << bc.timestamp() << " with magnetic field of " << d_bz << " kG"; + + // initialize MLResponse + if (fElectronCut.usePIDML) { + static constexpr int nClassesMl = 2; + const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + const std::vector labelsClasses = {"Background", "Signal"}; + const uint32_t nBinsMl = fElectronCut.binsMl.value.size() - 1; + const std::vector labelsBins(nBinsMl, "bin"); + double cutsMlArr[nBinsMl][nClassesMl]; + for (uint32_t i = 0; i < nBinsMl; i++) { + cutsMlArr[i][0] = 0.0; + cutsMlArr[i][1] = fElectronCut.cutsMl.value[i]; + } + o2::framework::LabeledArray cutsMltmp = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + mlResponsePID.configure(fElectronCut.binsMl.value, cutsMltmp, cutDirMl, nClassesMl); + if (fElectronCut.loadModelsFromCCDB) { + mlResponsePID.setModelPathsCCDB(fElectronCut.onnxFileNames, ccdbApi, fElectronCut.onnxPathsCCDB, bc.timestamp()); + } else { + mlResponsePID.setModelPathsLocal(fElectronCut.onnxFileNames); + } + mlResponsePID.cacheInputFeaturesIndices(fElectronCut.namesInputFeatures); + mlResponsePID.cacheBinningIndex(fElectronCut.nameBinningFeature); + mlResponsePID.init(fElectronCut.enableOptimizations); + } // end of ML PID + + if (fConfigDFeT.useML) { + static constexpr int nClassesMl = 4; + const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + const std::vector labelsClasses = {"prompt", "prompthc", "nonprompthc", "hb"}; + const uint32_t nBinsMl = fConfigDFeT.binsMl.value.size() - 1; + const std::vector labelsBins(nBinsMl, "bin"); + double cutsMlArr[nBinsMl][nClassesMl]; + for (uint32_t i = 0; i < nBinsMl; i++) { + cutsMlArr[i][0] = 0.0; + cutsMlArr[i][1] = 0.0; + cutsMlArr[i][2] = 0.0; + cutsMlArr[i][3] = 0.0; + } + o2::framework::LabeledArray cutsMltmp = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + mlResponseSCTeT.configure(fConfigDFeT.binsMl.value, cutsMltmp, cutDirMl, nClassesMl); + if (fConfigDFeT.loadModelsFromCCDB) { + mlResponseSCTeT.setModelPathsCCDB(fConfigDFeT.onnxFileNames, ccdbApi, fConfigDFeT.onnxPathsCCDB, bc.timestamp()); + } else { + mlResponseSCTeT.setModelPathsLocal(fConfigDFeT.onnxFileNames); + } + mlResponseSCTeT.cacheInputFeaturesIndices(fConfigDFeT.namesInputFeatures); + mlResponseSCTeT.cacheBinningIndex(fConfigDFeT.nameBinningFeature); + mlResponseSCTeT.init(fConfigDFeT.enableOptimizations); + } // end of ML SCTeT + + if (fConfigDFeV0.useML) { + static constexpr int nClassesMl = 4; + const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + const std::vector labelsClasses = {"prompt", "prompthc", "nonprompthc", "hb"}; + const uint32_t nBinsMl = fConfigDFeV0.binsMl.value.size() - 1; + const std::vector labelsBins(nBinsMl, "bin"); + double cutsMlArr[nBinsMl][nClassesMl]; + for (uint32_t i = 0; i < nBinsMl; i++) { + cutsMlArr[i][0] = 0.0; + cutsMlArr[i][1] = 0.0; + cutsMlArr[i][2] = 0.0; + cutsMlArr[i][3] = 0.0; + } + o2::framework::LabeledArray cutsMltmp = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + mlResponseSCTeV0.configure(fConfigDFeV0.binsMl.value, cutsMltmp, cutDirMl, nClassesMl); + if (fConfigDFeV0.loadModelsFromCCDB) { + mlResponseSCTeV0.setModelPathsCCDB(fConfigDFeV0.onnxFileNames, ccdbApi, fConfigDFeV0.onnxPathsCCDB, bc.timestamp()); + } else { + mlResponseSCTeV0.setModelPathsLocal(fConfigDFeV0.onnxFileNames); + } + mlResponseSCTeV0.cacheInputFeaturesIndices(fConfigDFeV0.namesInputFeatures); + mlResponseSCTeV0.cacheBinningIndex(fConfigDFeV0.nameBinningFeature); + mlResponseSCTeV0.init(fConfigDFeV0.enableOptimizations); + } // end of ML SCTeV0 + + if (fConfigDFeC.useML) { + static constexpr int nClassesMl = 4; + const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + const std::vector labelsClasses = {"prompt", "prompthc", "nonprompthc", "hb"}; + const uint32_t nBinsMl = fConfigDFeC.binsMl.value.size() - 1; + const std::vector labelsBins(nBinsMl, "bin"); + double cutsMlArr[nBinsMl][nClassesMl]; + for (uint32_t i = 0; i < nBinsMl; i++) { + cutsMlArr[i][0] = 0.0; + cutsMlArr[i][1] = 0.0; + cutsMlArr[i][2] = 0.0; + cutsMlArr[i][3] = 0.0; + } + o2::framework::LabeledArray cutsMltmp = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + mlResponseSCTeC.configure(fConfigDFeC.binsMl.value, cutsMltmp, cutDirMl, nClassesMl); + if (fConfigDFeC.loadModelsFromCCDB) { + mlResponseSCTeC.setModelPathsCCDB(fConfigDFeC.onnxFileNames, ccdbApi, fConfigDFeC.onnxPathsCCDB, bc.timestamp()); + } else { + mlResponseSCTeC.setModelPathsLocal(fConfigDFeC.onnxFileNames); + } + mlResponseSCTeC.cacheInputFeaturesIndices(fConfigDFeC.namesInputFeatures); + mlResponseSCTeC.cacheBinningIndex(fConfigDFeC.nameBinningFeature); + mlResponseSCTeC.init(fConfigDFeC.enableOptimizations); + } // end of ML SCTeC + + mRunNumber = bc.runNumber(); + mTOFResponse->processSetup(bc); + } + + template + void addHistograms(THistoregistry& registry) + { + registry.add("Track/hPt", "pT;p_{T} (GeV/c)", o2::framework::HistType::kTH1F, {{1000, 0.0f, 10}}, false); + registry.add("Track/hQoverPt", "q/pT;q/p_{T} (GeV/c)^{-1}", o2::framework::HistType::kTH1F, {{4000, -20, 20}}, false); + registry.add("Track/hEtaPhi", "#eta vs. #varphi;#varphi (rad.);#eta", o2::framework::HistType::kTH2F, {{180, 0, 2 * M_PI}, {20, -1.0f, 1.0f}}, false); + registry.add("Track/hDCAxyz", "DCA xy vs. z;DCA_{xy} (cm);DCA_{z} (cm)", o2::framework::HistType::kTH2F, {{200, -1.0f, 1.0f}, {200, -1.0f, 1.0f}}, false); + registry.add("Track/hDCAxyzSigma", "DCA xy vs. z;DCA_{xy} (#sigma);DCA_{z} (#sigma)", o2::framework::HistType::kTH2F, {{200, -10.0f, 10.0f}, {200, -10.0f, 10.0f}}, false); + registry.add("Track/hDCAxyRes_Pt", "DCA_{xy} resolution vs. pT;p_{T} (GeV/c);DCA_{xy} resolution (#mum)", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {500, 0., 500}}, false); + registry.add("Track/hDCAzRes_Pt", "DCA_{z} resolution vs. pT;p_{T} (GeV/c);DCA_{z} resolution (#mum)", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {500, 0., 500}}, false); + registry.add("Track/hNclsTPC", "number of TPC clusters", o2::framework::HistType::kTH1F, {{161, -0.5, 160.5}}, false); + registry.add("Track/hNcrTPC", "number of TPC crossed rows", o2::framework::HistType::kTH1F, {{161, -0.5, 160.5}}, false); + registry.add("Track/hChi2TPC", "chi2/number of TPC clusters", o2::framework::HistType::kTH1F, {{100, 0, 10}}, false); + registry.add("Track/hChi2TOF", "chi2 of TOF", o2::framework::HistType::kTH1F, {{100, 0, 10}}, false); + registry.add("Track/hTPCNcr2Nf", "TPC Ncr/Nfindable", o2::framework::HistType::kTH1F, {{200, 0, 2}}, false); + registry.add("Track/hTPCNcls2Nf", "TPC Ncls/Nfindable", o2::framework::HistType::kTH1F, {{200, 0, 2}}, false); + registry.add("Track/hTPCNclsShared", "TPC Ncls shared/Ncls;p_{T} (GeV/c);N_{cls}^{shared}/N_{cls} in TPC", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, 0, 1}}, false); + registry.add("Track/hNclsITS", "number of ITS clusters", o2::framework::HistType::kTH1F, {{8, -0.5, 7.5}}, false); + registry.add("Track/hChi2ITS", "chi2/number of ITS clusters", o2::framework::HistType::kTH1F, {{100, 0, 10}}, false); + registry.add("Track/hITSClusterMap", "ITS cluster map", o2::framework::HistType::kTH1F, {{128, -0.5, 127.5}}, false); + registry.add("Track/hTPCdEdx", "TPC dE/dx;p_{in} (GeV/c);TPC dE/dx (a.u.)", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {200, 0, 200}}, false); + registry.add("Track/hTPCdEdxMC", "TPC dE/dx;p_{in} (GeV/c);TPC dE/dx (a.u.)", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {200, 0, 200}}, false); + registry.add("Track/hTPCNsigmaEl", "TPC n sigma el;p_{in} (GeV/c);n #sigma_{e}^{TPC}", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); + registry.add("Track/hTPCNsigmaPi", "TPC n sigma pi;p_{in} (GeV/c);n #sigma_{#pi}^{TPC}", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); + registry.add("Track/hTPCNsigmaKa", "TPC n sigma ka;p_{in} (GeV/c);n #sigma_{K}^{TPC}", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); + registry.add("Track/hTPCNsigmaPr", "TPC n sigma pr;p_{in} (GeV/c);n #sigma_{p}^{TPC}", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); + registry.add("Track/hTOFbeta", "TOF beta;p_{pv} (GeV/c);#beta", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {240, 0, 1.2}}, false); + registry.add("Track/hTOFNsigmaEl", "TOF n sigma el;p_{in} (GeV/c);n #sigma_{e}^{TOF}", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); + registry.add("Track/hTOFNsigmaPi", "TOF n sigma pi;p_{in} (GeV/c);n #sigma_{#pi}^{TOF}", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); + registry.add("Track/hTOFNsigmaKa", "TOF n sigma ka;p_{in} (GeV/c);n #sigma_{K}^{TOF}", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); + registry.add("Track/hTOFNsigmaPr", "TOF n sigma pr;p_{in} (GeV/c);n #sigma_{p}^{TOF}", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); + registry.add("Track/hMeanClusterSizeITS", "mean cluster size ITS;p_{pv} (GeV/c); #times cos(#lambda)", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {150, 0, 15}}, false); + registry.add("Track/hMeanClusterSizeITSib", "mean cluster size ITSib;p_{pv} (GeV/c); #times cos(#lambda)", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {150, 0, 15}}, false); + registry.add("Track/hMeanClusterSizeITSob", "mean cluster size ITSob;p_{pv} (GeV/c); #times cos(#lambda)", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {150, 0, 15}}, false); + registry.add("Track/hProbElBDT", "probability to be e from BDT;p_{in} (GeV/c);BDT score;", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {100, 0, 1}}, false); + registry.add("Track/hNe", "electron counts;N_{e} per collision", o2::framework::HistType::kTH1F, {{51, -0.5, 50.5}}, false); + + registry.add("Prefilter/Track/hPt", "pT;p_{T} (GeV/c)", o2::framework::HistType::kTH1F, {{1000, 0.0f, 10}}, false); + registry.add("Prefilter/Track/hEtaPhi", "#eta vs. #varphi;#varphi (rad.);#eta", o2::framework::HistType::kTH2F, {{180, 0, 2 * M_PI}, {40, -2, 2}}, false); + registry.add("Prefilter/Track/hTPCdEdx", "TPC dE/dx vs. pin;p_{in} (GeV/c);TPC dE/dx", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {200, 0, 200}}, false); + registry.add("Prefilter/Track/hTOFbeta", "TOF #beta vs. p;p_{pv} (GeV/c);TOF #beta", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {600, 0, 1.2}}, false); + registry.add("Prefilter/Pair/hMvsPhiV", "mass vs. phiv;#varphi_{V} (rad.);m_{ee} (GeV/c^{2})", o2::framework::HistType::kTH2F, {{90, 0.f, M_PI}, {100, 0, 0.1}}); + + // for charged hadron + registry.add("SCT/Track/hs", "hs;p_{T} (GeV/c);#eta;#varphi (rad.)", o2::framework::HistType::kTHnSparseF, {{100, 0, 10}, {80, -2, 2}, {36, 0, 2 * M_PI}}, false); + registry.add("SCT/Track/hDCA", "DCA xy vs. z;DCA_{xy} (cm);DCA_{z} (cm)", o2::framework::HistType::kTH2F, {{200, -1, 1}, {200, -1, 1}}, false); + registry.add("SCT/Track/hTPCdEdx", "TPC dE/dx vs. pin;p_{in} (GeV/c);TPC dE/dx", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {200, 0, 200}}, false); + registry.add("SCT/Track/hTOFbeta", "TOF #beta vs. p;p_{pv} (GeV/c);TOF #beta", o2::framework::HistType::kTH2F, {{1000, 0, 10}, {600, 0, 1.2}}, false); + + // for V0 + registry.add("SCT/V0/hPt", "pT of V0;p_{T} (GeV/c)", o2::framework::HistType::kTH1F, {{100, 0, 10}}, false); + registry.add("SCT/V0/hYPhi_K0S", "rapidity vs. #varphi of V0;#varphi (rad.);rapidity_{K0S}", o2::framework::HistType::kTH2F, {{90, 0, 2 * M_PI}, {80, -2, +2}}, false); + registry.add("SCT/V0/hYPhi_Lambda", "rapidity vs. #varphi of V0;#varphi (rad.);rapidity_{#Lambda}", o2::framework::HistType::kTH2F, {{90, 0, 2 * M_PI}, {80, -2, +2}}, false); + registry.add("SCT/V0/hAP", "AP plot;#alpha;q_{T} (GeV/c)", o2::framework::HistType::kTH2F, {{200, -1, 1}, {250, 0, 0.25}}, false); + registry.add("SCT/V0/hLxy", "decay length from PV;L_{xy} (cm)", o2::framework::HistType::kTH1F, {{100, 0, 10}}, false); + registry.add("SCT/V0/hCosPA", "cosPA;cosine of pointing angle", o2::framework::HistType::kTH1F, {{100, 0.9, 1}}, false); + registry.add("SCT/V0/hDCA2Legs", "distance between 2 legs at PCA;distance between 2 legs (cm)", o2::framework::HistType::kTH1F, {{100, 0, 1}}, false); + registry.add("SCT/V0/hMassK0S", "K0S mass;m_{#pi#pi} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 0.45, 0.55}}, false); + registry.add("SCT/V0/hMassLambda", "Lambda mass;m_{p#pi^{#minus}} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 1.08, 1.18}}, false); + registry.add("SCT/V0/hMassAntiLambda", "Anti-Lambda mass;m_{#bar{p}#pi^{+}} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 1.08, 1.18}}, false); + registry.add("SCT/V0/hMassGamma_misid", "#gamma mass;m_{ee} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 0, 0.1}}, false); + registry.add("SCT/V0/hMassK0S_misid", "K0S mass;m_{#pi#pi} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 0.45, 0.55}}, false); + registry.add("SCT/V0/hMassLambda_misid", "Lambda mass;m_{p#pi^{#minus}} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 1.08, 1.18}}, false); + registry.add("SCT/V0/hMassAntiLambda_misid", "Anti-Lambda mass;m_{#bar{p}#pi^{+}} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 1.08, 1.18}}, false); + + // for cascadeSCT/ + registry.add("SCT/Cascade/hPt", "pT of cascade;p_{T} (GeV/c)", o2::framework::HistType::kTH1F, {{100, 0, 10}}, false); + registry.add("SCT/Cascade/hYPhi_Xi", "rapidity vs. #varphi of cascade;#varphi (rad.);rapidity_{#Xi}", o2::framework::HistType::kTH2F, {{90, 0, 2 * M_PI}, {80, -2, +2}}, false); + registry.add("SCT/Cascade/hYPhi_Omega", "rapidity vs. #varphi of cascade;#varphi (rad.);rapidity_{#Omega}", o2::framework::HistType::kTH2F, {{90, 0, 2 * M_PI}, {80, -2, +2}}, false); + registry.add("SCT/Cascade/hCosPA", "cosPA;cosine of pointing angle", o2::framework::HistType::kTH1F, {{100, 0.9, 1}}, false); + registry.add("SCT/Cascade/hDCA2Legs", "distance between 2 legs at PCA;distance between 2 legs (cm)", o2::framework::HistType::kTH1F, {{100, 0, 1}}, false); + registry.add("SCT/Cascade/hV0CosPA", "cosPA of V0 in cascade;cosine of pointing angle", o2::framework::HistType::kTH1F, {{100, 0.9, 1}}, false); + registry.add("SCT/Cascade/hV0DCA2Legs", "distance between 2 legs at PCA of V0 in cascade;distance between 2 legs (cm)", o2::framework::HistType::kTH1F, {{100, 0, 1}}, false); + registry.add("SCT/Cascade/hMassLambda", "Lambda mass;m_{p#pi^{-}} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 1.08, 1.18}}, false); + registry.add("SCT/Cascade/hMassXi", "#Xi mass;m_{#Lambda#pi} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 1.27, 1.37}}, false); + registry.add("SCT/Cascade/hMassOmega", "#Omega mass;m_{#LambdaK} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 1.62, 1.72}}, false); + registry.add("SCT/Cascade/hMassXi_misid", "#Xi mass;m_{#Lambda#pi} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 1.27, 1.37}}, false); + registry.add("SCT/Cascade/hMassOmega_misid", "#Omega mass;m_{#LambdaK} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 1.62, 1.72}}, false); + + registry.add("SCT/eT/hImpactParameter", "eH impact parameter;IP in XY (cm);IP in Z (cm);", o2::framework::HistType::kTH2F, {{200, -1, 1}, {200, -1, 1}}, false); + registry.add("SCT/eT/hDecayLength", "eH decay length;L_{xy} (cm);L_{z} (cm);", o2::framework::HistType::kTH2F, {{100, 0, 1}, {200, -1, 1}}, false); + registry.add("SCT/eT/hCosPA", "eH cosPA;cosine of pointing angle", o2::framework::HistType::kTH1F, {{200, -1, 1}}, false); + registry.add("SCT/eT/hDCA2legs", "dca between e and hadron;dca between 2legs (cm)", o2::framework::HistType::kTH1F, {{100, 0, 1}}, false); + registry.add("SCT/eT/hMass", "invariant mass of e-h;m_{eh} (GeV/c^{2})", o2::framework::HistType::kTH1F, {{100, 0, 10}}, false); + registry.addClone("SCT/eT/", "SCT/eV0/"); + registry.addClone("SCT/eT/", "SCT/eC/"); + } + + template + void calculateTOFNSigmaWithReassociation(TCollisions const& collisions, TBCs const&, TTracks const& tracks, TTrackAssoc const& trackIndices, TSliceCache const&, TPresliceTrack const& perColTrack, TPresliceTrackAssoc const& trackIndicesPerCollision) + { + if (fElectronCut.useTOFNSigmaDeltaBC) { + if constexpr (withTTCA) { + std::unordered_map mapCollisionTime; + std::unordered_map mapCollisionTimeError; + for (const auto& track : tracks) { + if (mapCollisionTime.find(track.collisionId()) == mapCollisionTime.end()) { + mapCollisionTime[track.collisionId()] = track.tofEvTime(); + mapCollisionTimeError[track.collisionId()] = track.tofEvTimeErr(); + } + } + + for (const auto& collision : collisions) { + if (mapCollisionTime.find(collision.globalIndex()) == mapCollisionTime.end()) { + continue; + } + auto bcCollision = collision.template bc_as(); + auto trackIdsThisCollision = trackIndices.sliceBy(trackIndicesPerCollision, collision.globalIndex()); + for (const auto& trackId : trackIdsThisCollision) { + auto track = trackId.template track_as(); + if (!track.hasITS() || !track.hasTPC()) { // apply only minimal cut + continue; + } + + if (track.hasTOF() && track.has_collision()) { // TTCA may use orphan tracks. + auto bcTrack = track.template collision_as().template bc_as(); + fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = mTOFResponse->nSigma(track.tofSignalInAnotherBC(bcTrack.globalBC(), bcCollision.globalBC()), track.tofExpMom(), track.length(), track.p(), track.eta(), mapCollisionTime[collision.globalIndex()], mapCollisionTimeError[collision.globalIndex()]); + fMapTOFNsigmaPiReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = mTOFResponse->nSigma(track.tofSignalInAnotherBC(bcTrack.globalBC(), bcCollision.globalBC()), track.tofExpMom(), track.length(), track.p(), track.eta(), mapCollisionTime[collision.globalIndex()], mapCollisionTimeError[collision.globalIndex()]); + fMapTOFNsigmaKaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = mTOFResponse->nSigma(track.tofSignalInAnotherBC(bcTrack.globalBC(), bcCollision.globalBC()), track.tofExpMom(), track.length(), track.p(), track.eta(), mapCollisionTime[collision.globalIndex()], mapCollisionTimeError[collision.globalIndex()]); + fMapTOFNsigmaPrReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = mTOFResponse->nSigma(track.tofSignalInAnotherBC(bcTrack.globalBC(), bcCollision.globalBC()), track.tofExpMom(), track.length(), track.p(), track.eta(), mapCollisionTime[collision.globalIndex()], mapCollisionTimeError[collision.globalIndex()]); + fMapTOFBetaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.length() / (track.tofSignalInAnotherBC(bcTrack.globalBC(), bcCollision.globalBC()) - mapCollisionTime[collision.globalIndex()]) / (TMath::C() * 1e+2 * 1e-12); + ; + } else { + fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaEl(); + fMapTOFNsigmaPiReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaPi(); + fMapTOFNsigmaKaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaKa(); + fMapTOFNsigmaPrReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaPr(); + fMapTOFBetaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.beta(); + } + } // end of track loop + } // end of collision loop + mapCollisionTime.clear(); + mapCollisionTimeError.clear(); + } else { // without TTCA + for (const auto& collision : collisions) { + auto tracks_per_coll = tracks.sliceBy(perColTrack, collision.globalIndex()); + for (const auto& track : tracks_per_coll) { + if (!track.hasITS() || !track.hasTPC()) { // apply only minimal cut + continue; + } + fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaEl(); + fMapTOFNsigmaPiReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaPi(); + fMapTOFNsigmaKaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaKa(); + fMapTOFNsigmaPrReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaPr(); + fMapTOFBetaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.beta(); + } + } // end of track loop + } // end of collision loop + } else { + if constexpr (withTTCA) { + for (const auto& collision : collisions) { + auto trackIdsThisCollision = trackIndices.sliceBy(trackIndicesPerCollision, collision.globalIndex()); + for (const auto& trackId : trackIdsThisCollision) { + auto track = trackId.template track_as(); + if (!track.hasITS() || !track.hasTPC()) { // apply only minimal cut + continue; + } + fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaEl(); + fMapTOFNsigmaPiReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaPi(); + fMapTOFNsigmaKaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaKa(); + fMapTOFNsigmaPrReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaPr(); + fMapTOFBetaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.beta(); + } // end of track loop + } // end of collision loop + } else { + for (const auto& collision : collisions) { + auto tracks_per_coll = tracks.sliceBy(perColTrack, collision.globalIndex()); + for (const auto& track : tracks_per_coll) { + if (!track.hasITS() || !track.hasTPC()) { // apply only minimal cut + continue; + } + fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaEl(); + fMapTOFNsigmaPiReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaPi(); + fMapTOFNsigmaKaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaKa(); + fMapTOFNsigmaPrReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.tofNSigmaPr(); + fMapTOFBetaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())] = track.beta(); + } + } // end of track loop + } // end of collision loop + } + } + + template + void fillMapMLPID(TCollision const& collision, TTrack const& track) + { + if (fElectronCut.usePIDML) { + o2::dataformats::DCA mDcaInfoCov; + mDcaInfoCov.set(999, 999, 999, 999, 999); + auto trackParCov = getTrackParCov(track); + trackParCov.setPID(o2::track::PID::Electron); + bool isPropOK = o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov); + if (!isPropOK) { + return; + } + + if (!isElectron_TOFif(track, collision, fElectronCut)) { // minimal n sigma cut is taken from the main electron cut. + return; + } + + o2::analysis::pwgem::dilepton::mlpid::candidate candidate; + candidate.tpcInnerParam = track.tpcInnerParam(); + candidate.tpcNClsFound = track.tpcNClsFound(); + candidate.tpcChi2NCl = track.tpcChi2NCl(); + candidate.tpcNSigmaEl = track.tpcNSigmaEl(); + candidate.tofNSigmaEl = fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + int total_cluster_size_ob = 0, nl_ob = 0; + for (unsigned int layer = 3; layer < 7; layer++) { + int cluster_size_per_layer = track.itsClsSizeInLayer(layer); + if (cluster_size_per_layer > 0) { + nl_ob++; + } + total_cluster_size_ob += cluster_size_per_layer; + } + candidate.meanClusterSizeITSobCosTgl = static_cast(total_cluster_size_ob) / static_cast(nl_ob) * std::cos(std::atan(trackParCov.getTgl())); + + std::vector inputFeatures = mlResponsePID.getInputFeatures(candidate); + float binningFeature = mlResponsePID.getBinningFeature(candidate); + + int pbin = lower_bound(fElectronCut.binsMl.value.begin(), fElectronCut.binsMl.value.end(), binningFeature) - fElectronCut.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(fElectronCut.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(fElectronCut.binsMl.value.size()) - 2; + } + + float probaEl = mlResponsePID.getModelOutput(inputFeatures, pbin)[1]; // 0: hadron, 1:electron + fMapProbaEl[std::make_pair(collision.globalIndex(), track.globalIndex())] = probaEl; + } else { + fMapProbaEl[std::make_pair(collision.globalIndex(), track.globalIndex())] = 1.0; + } + } + + template + bool checkTrack(TCollision const& collision, TTrack const& track) + { + if constexpr (isMC) { + if (!track.has_mcParticle()) { + return false; + } + if (fElectronCut.storeOnlyTrueElectronMC) { + const auto& mcParticle = track.template mcParticle_as(); + if (std::abs(mcParticle.pdgCode()) != 11) { + return false; + } + } + } + + float tofNSigmaEl = fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + if (fElectronCut.requireTOF && !(track.hasTOF() && std::fabs(tofNSigmaEl) < fElectronCut.maxTOFNsigmaEl)) { + return false; + } + + if (!track.hasITS()) { + return false; + } + + if (track.itsChi2NCl() < fElectronCut.minchi2its || fElectronCut.maxchi2its < track.itsChi2NCl()) { + return false; + } + if (track.itsNCls() < fElectronCut.min_ncluster_its) { + return false; + } + if (track.itsNClsInnerBarrel() < fElectronCut.min_ncluster_itsib) { + return false; + } + + if (!fElectronCut.includeITSsa && (!track.hasITS() || !track.hasTPC())) { + return false; + } + + if (track.hasTPC()) { + if (track.tpcChi2NCl() < fElectronCut.minchi2tpc || fElectronCut.maxchi2tpc < track.tpcChi2NCl()) { + return false; + } + + if (track.tpcNClsFound() < fElectronCut.min_ncluster_tpc) { + return false; + } + + if (track.tpcNClsCrossedRows() < fElectronCut.mincrossedrows) { + return false; + } + + if (track.tpcCrossedRowsOverFindableCls() < fElectronCut.min_tpc_cr_findable_ratio) { + return false; + } + + if (track.tpcFractionSharedCls() > fElectronCut.max_frac_shared_clusters_tpc) { + return false; + } + } + + o2::dataformats::DCA mDcaInfoCov; + mDcaInfoCov.set(999, 999, 999, 999, 999); + auto trackParCov = getTrackParCov(track); + trackParCov.setPID(o2::track::PID::Electron); + bool isPropOK = o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov); + if (!isPropOK) { + return false; + } + float dcaXY = mDcaInfoCov.getY(); + float dcaZ = mDcaInfoCov.getZ(); + + if (std::fabs(dcaXY) > fElectronCut.dca_xy_max || std::fabs(dcaZ) > fElectronCut.dca_z_max) { + return false; + } + + if (trackParCov.getPt() < fElectronCut.minpt || std::fabs(trackParCov.getEta()) > fElectronCut.maxeta) { + return false; + } + + int total_cluster_size = 0, nl = 0; + for (unsigned int layer = 0; layer < 7; layer++) { + int cluster_size_per_layer = track.itsClsSizeInLayer(layer); + if (cluster_size_per_layer > 0) { + nl++; + } + total_cluster_size += cluster_size_per_layer; + } + + if (fElectronCut.maxMeanITSClusterSize < static_cast(total_cluster_size) / static_cast(nl) * std::cos(std::atan(trackParCov.getTgl()))) { + return false; + } + + return true; + } + + template + bool checkTrackPF(TCollision const&, TTrack const& track) + { + if constexpr (isMC) { + if (!track.has_mcParticle()) { + return false; + } + } + + if (!track.hasITS() || !track.hasTPC()) { + return false; + } + + if (track.itsChi2NCl() < fElectronPFCut.minchi2its || fElectronPFCut.maxchi2its < track.itsChi2NCl()) { + return false; + } + if (track.itsNCls() < fElectronPFCut.min_ncluster_its) { + return false; + } + if (track.itsNClsInnerBarrel() < fElectronPFCut.min_ncluster_itsib) { + return false; + } + + if (track.tpcChi2NCl() < fElectronPFCut.minchi2tpc || fElectronPFCut.maxchi2tpc < track.tpcChi2NCl()) { + return false; + } + + if (track.tpcNClsFound() < fElectronPFCut.min_ncluster_tpc) { + return false; + } + + if (track.tpcNClsCrossedRows() < fElectronPFCut.mincrossedrows) { + return false; + } + + if (track.tpcCrossedRowsOverFindableCls() < fElectronPFCut.min_tpc_cr_findable_ratio) { + return false; + } + + if (track.tpcFractionSharedCls() > fElectronPFCut.max_frac_shared_clusters_tpc) { + return false; + } + + o2::dataformats::DCA mDcaInfoCov; + mDcaInfoCov.set(999, 999, 999, 999, 999); + auto trackParCov = getTrackParCov(track); + trackParCov.setPID(o2::track::PID::Electron); + bool isPropOK = o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov); + if (!isPropOK) { + return false; + } + float dcaXY = mDcaInfoCov.getY(); + float dcaZ = mDcaInfoCov.getZ(); + + if (std::fabs(dcaXY) > fElectronPFCut.dca_xy_max || std::fabs(dcaZ) > fElectronPFCut.dca_z_max) { + return false; + } + + if (trackParCov.getPt() < fElectronPFCut.minpt || std::fabs(trackParCov.getEta()) > fElectronPFCut.maxeta) { + return false; + } + + int total_cluster_size = 0, nl = 0; + for (unsigned int layer = 0; layer < 7; layer++) { + int cluster_size_per_layer = track.itsClsSizeInLayer(layer); + if (cluster_size_per_layer > 0) { + nl++; + } + total_cluster_size += cluster_size_per_layer; + } + + if (fElectronPFCut.maxMeanITSClusterSize < static_cast(total_cluster_size) / static_cast(nl) * std::cos(std::atan(trackParCov.getTgl()))) { + return false; + } + + return true; + } + + template + bool isElectron(TCollision const& collision, TTrack const& track, TConfig const& eCut) + { + if (fElectronCut.usePIDML) { // keep fElectronCut here intentinoally + if (isElectron_TOFif(track, collision, eCut)) { + o2::analysis::pwgem::dilepton::mlpid::candidate candidate; + candidate.tpcInnerParam = track.tpcInnerParam(); + float binningFeature = mlResponsePID.getBinningFeature(candidate); + int pbin = lower_bound(eCut.binsMl.value.begin(), eCut.binsMl.value.end(), binningFeature) - eCut.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(eCut.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(eCut.binsMl.value.size()) - 2; + } + return fMapProbaEl[std::make_pair(collision.globalIndex(), track.globalIndex())] > eCut.cutsMl.value[pbin]; + } else { + return false; + } + } else { + return isElectron_TPChadrej(track, collision, eCut) || isElectron_TOFreq(track, collision, eCut); + } + } + + template + bool isElectron_TOFif(TTrack const& track, TCollision const& collision, TConfig const& eCut) + { + float tofNSigmaEl = fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + bool is_EL_TPC = eCut.minTPCNsigmaEl.value < track.tpcNSigmaEl() && track.tpcNSigmaEl() < eCut.maxTPCNsigmaEl.value; + bool is_EL_TOF = track.hasTOF() ? (std::fabs(tofNSigmaEl) < eCut.maxTOFNsigmaEl.value) : true; // TOFif + return is_EL_TPC && is_EL_TOF; + } + + template + bool isElectron_TPChadrej(TTrack const& track, TCollision const& collision, TConfig const& eCut) + { + float tofNSigmaEl = fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + if (track.tpcNSigmaEl() < eCut.minTPCNsigmaEl.value || eCut.maxTPCNsigmaEl.value < track.tpcNSigmaEl()) { + return false; + } + if (eCut.minTPCNsigmaPi.value < track.tpcNSigmaPi() && track.tpcNSigmaPi() < eCut.maxTPCNsigmaPi.value && track.tpcInnerParam() < eCut.max_pin_for_pion_rejection.value) { + return false; + } + if (eCut.minTPCNsigmaKa.value < track.tpcNSigmaKa() && track.tpcNSigmaKa() < eCut.maxTPCNsigmaKa.value) { + return false; + } + if (eCut.minTPCNsigmaPr.value < track.tpcNSigmaPr() && track.tpcNSigmaPr() < eCut.maxTPCNsigmaPr.value) { + return false; + } + if (track.hasTOF() && (eCut.maxTOFNsigmaEl.value < std::fabs(tofNSigmaEl))) { + return false; + } + return true; + } + + template + bool isElectron_TOFreq(TTrack const& track, TCollision const& collision, TConfig const& eCut) + { + float tofNSigmaEl = fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + if (eCut.minTPCNsigmaPi.value < track.tpcNSigmaPi() && track.tpcNSigmaPi() < eCut.maxTPCNsigmaPi.value && (eCut.min_pin_for_pion_rejection.value < track.tpcInnerParam() && track.tpcInnerParam() < eCut.max_pin_for_pion_rejection.value)) { + return false; + } + return eCut.minTPCNsigmaEl.value < track.tpcNSigmaEl() && track.tpcNSigmaEl() < eCut.maxTPCNsigmaEl.value && std::fabs(tofNSigmaEl) < eCut.maxTOFNsigmaEl.value; + } + + template + bool isSelectedHadron(TCollision const&, TTrack const& track, TTrackParCov const& trackParCov, const float dcaXY, const float dcaZ) + { + if (!track.hasITS() || !track.hasTPC()) { + return false; + } + + if (trackParCov.getPt() < fHadronCut.cfg_min_pt_track || fHadronCut.cfg_max_pt_track < trackParCov.getPt()) { + return false; + } + + if (trackParCov.getEta() < fHadronCut.cfg_min_eta_track || fHadronCut.cfg_max_eta_track < trackParCov.getEta()) { + return false; + } + + if (std::fabs(dcaXY) > fHadronCut.cfg_max_dcaxy) { + return false; + } + + if (std::fabs(dcaZ) > fHadronCut.cfg_max_dcaz) { + return false; + } + + if (track.itsChi2NCl() < 0.f || fHadronCut.cfg_max_chi2its < track.itsChi2NCl()) { + return false; + } + + if (track.itsNCls() < fHadronCut.cfg_min_ncluster_its) { + return false; + } + + if (track.itsNClsInnerBarrel() < fHadronCut.cfg_min_ncluster_itsib) { + return false; + } + + if (track.tpcChi2NCl() < 0.f || fHadronCut.cfg_max_chi2tpc < track.tpcChi2NCl()) { + return false; + } + + if (track.tpcNClsFound() < fHadronCut.cfg_min_ncluster_tpc) { + return false; + } + + if (track.tpcNClsCrossedRows() < fHadronCut.cfg_min_ncrossedrows_tpc) { + return false; + } + + if (track.tpcCrossedRowsOverFindableCls() < fHadronCut.cfg_min_cr2findable_ratio_tpc) { + return false; + } + + if (track.tpcFractionSharedCls() > fHadronCut.cfg_max_frac_shared_clusters_tpc) { + return false; + } + + if (fHadronCut.requirePiKaPr && !isPiKaPr(track)) { + return false; + } + + return true; + } + + template + void fillElectronTable(TCollision const& collision, TTrack const& track, TTrackParCov const& trackParCov, const float dcaXY, const float dcaZ, TProducts& products, THistoregistry& registry) + { + float pt = trackParCov.getPt(); + float eta = trackParCov.getEta(); + float phi = RecoDecay::constrainAngle(trackParCov.getPhi(), 0, 1U); + + float tofNSigmaEl = fMapTOFNsigmaElReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + float tofNSigmaPi = fMapTOFNsigmaPiReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + float tofNSigmaKa = fMapTOFNsigmaKaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + float tofNSigmaPr = fMapTOFNsigmaPrReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + float beta = fMapTOFBetaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]; + float probaEl = fMapProbaEl[std::make_pair(collision.globalIndex(), track.globalIndex())]; + + bool isAssociatedToMPC = collision.globalIndex() == track.collisionId(); + float mcTunedTPCSignal = 0.f; + if constexpr (isMC) { + mcTunedTPCSignal = track.mcTunedTPCSignal(); + } + + products.electronTable(collision.globalIndex(), track.globalIndex(), track.sign(), + pt, eta, phi, + dcaXY, dcaZ, trackParCov.getSigmaY2(), trackParCov.getSigmaZY(), trackParCov.getSigmaZ2(), + track.tpcNClsFindable(), track.tpcNClsFindableMinusFound(), track.tpcNClsFindableMinusPID(), track.tpcNClsFindableMinusCrossedRows(), track.tpcNClsShared(), + track.tpcChi2NCl(), track.tpcInnerParam(), + track.tpcSignal(), track.tpcNSigmaEl(), track.tpcNSigmaPi(), track.tpcNSigmaKa(), track.tpcNSigmaPr(), + beta, tofNSigmaEl, + track.itsClusterSizes(), + track.itsChi2NCl(), track.tofChi2(), track.detectorMap(), + isAssociatedToMPC, false, probaEl, track.flags(), mcTunedTPCSignal); + + products.electronCovTable( + trackParCov.getX(), + trackParCov.getAlpha(), + trackParCov.getY(), + trackParCov.getZ(), + trackParCov.getSnp(), + // trackParCov.getTgl(), + // trackParCov.getSigmaY2(), + // trackParCov.getSigmaZY(), + // trackParCov.getSigmaZ2(), + trackParCov.getSigmaSnpY(), + trackParCov.getSigmaSnpZ(), + trackParCov.getSigmaSnp2(), + trackParCov.getSigmaTglY(), + trackParCov.getSigmaTglZ(), + trackParCov.getSigmaTglSnp(), + trackParCov.getSigmaTgl2(), + trackParCov.getSigma1PtY(), + trackParCov.getSigma1PtZ(), + trackParCov.getSigma1PtSnp(), + trackParCov.getSigma1PtTgl(), + trackParCov.getSigma1Pt2()); + + int total_cluster_size = 0, nl = 0; + for (unsigned int layer = 0; layer < 7; layer++) { + int cluster_size_per_layer = track.itsClsSizeInLayer(layer); + if (cluster_size_per_layer > 0) { + nl++; + } + total_cluster_size += cluster_size_per_layer; + } + + int total_cluster_size_ib = 0, nl_ib = 0; + for (unsigned int layer = 0; layer < 3; layer++) { + int cluster_size_per_layer = track.itsClsSizeInLayer(layer); + if (cluster_size_per_layer > 0) { + nl_ib++; + } + total_cluster_size_ib += cluster_size_per_layer; + } + + int total_cluster_size_ob = 0, nl_ob = 0; + for (unsigned int layer = 3; layer < 7; layer++) { + int cluster_size_per_layer = track.itsClsSizeInLayer(layer); + if (cluster_size_per_layer > 0) { + nl_ob++; + } + total_cluster_size_ob += cluster_size_per_layer; + } + + registry.fill(HIST("Track/hPt"), pt); + registry.fill(HIST("Track/hQoverPt"), track.sign() / pt); + registry.fill(HIST("Track/hEtaPhi"), phi, eta); + registry.fill(HIST("Track/hDCAxyz"), dcaXY, dcaZ); + registry.fill(HIST("Track/hDCAxyzSigma"), dcaXY / std::sqrt(trackParCov.getSigmaY2()), dcaZ / std::sqrt(trackParCov.getSigmaZ2())); + registry.fill(HIST("Track/hDCAxyRes_Pt"), pt, std::sqrt(trackParCov.getSigmaY2()) * 1e+4); // convert cm to um + registry.fill(HIST("Track/hDCAzRes_Pt"), pt, std::sqrt(trackParCov.getSigmaZ2()) * 1e+4); // convert cm to um + registry.fill(HIST("Track/hNclsITS"), track.itsNCls()); + registry.fill(HIST("Track/hNclsTPC"), track.tpcNClsFound()); + registry.fill(HIST("Track/hNcrTPC"), track.tpcNClsCrossedRows()); + registry.fill(HIST("Track/hTPCNcr2Nf"), track.tpcCrossedRowsOverFindableCls()); + registry.fill(HIST("Track/hTPCNcls2Nf"), track.tpcFoundOverFindableCls()); + registry.fill(HIST("Track/hTPCNclsShared"), track.pt(), track.tpcFractionSharedCls()); + registry.fill(HIST("Track/hChi2TPC"), track.tpcChi2NCl()); + registry.fill(HIST("Track/hChi2ITS"), track.itsChi2NCl()); + registry.fill(HIST("Track/hChi2TOF"), track.tofChi2()); + registry.fill(HIST("Track/hITSClusterMap"), track.itsClusterMap()); + registry.fill(HIST("Track/hTPCdEdx"), track.tpcInnerParam(), track.tpcSignal()); + registry.fill(HIST("Track/hTPCdEdxMC"), track.tpcInnerParam(), mcTunedTPCSignal); + registry.fill(HIST("Track/hTPCNsigmaEl"), track.tpcInnerParam(), track.tpcNSigmaEl()); + registry.fill(HIST("Track/hTPCNsigmaPi"), track.tpcInnerParam(), track.tpcNSigmaPi()); + registry.fill(HIST("Track/hTPCNsigmaKa"), track.tpcInnerParam(), track.tpcNSigmaKa()); + registry.fill(HIST("Track/hTPCNsigmaPr"), track.tpcInnerParam(), track.tpcNSigmaPr()); + registry.fill(HIST("Track/hTOFbeta"), trackParCov.getP(), beta); + registry.fill(HIST("Track/hTOFNsigmaEl"), track.tpcInnerParam(), tofNSigmaEl); + registry.fill(HIST("Track/hTOFNsigmaPi"), track.tpcInnerParam(), tofNSigmaPi); + registry.fill(HIST("Track/hTOFNsigmaKa"), track.tpcInnerParam(), tofNSigmaKa); + registry.fill(HIST("Track/hTOFNsigmaPr"), track.tpcInnerParam(), tofNSigmaPr); + registry.fill(HIST("Track/hMeanClusterSizeITS"), trackParCov.getP(), static_cast(total_cluster_size) / static_cast(nl) * std::cos(std::atan(trackParCov.getTgl()))); + registry.fill(HIST("Track/hMeanClusterSizeITSib"), trackParCov.getP(), static_cast(total_cluster_size_ib) / static_cast(nl_ib) * std::cos(std::atan(trackParCov.getTgl()))); + registry.fill(HIST("Track/hMeanClusterSizeITSob"), trackParCov.getP(), static_cast(total_cluster_size_ob) / static_cast(nl_ob) * std::cos(std::atan(trackParCov.getTgl()))); + registry.fill(HIST("Track/hProbElBDT"), track.tpcInnerParam(), probaEl); + } + + template + bool isPiKaPr(TTrack const& track) + { + bool is_pi_included_TPC = fHadronCut.cfg_min_TPCNsigmaPi < track.tpcNSigmaPi() && track.tpcNSigmaPi() < fHadronCut.cfg_max_TPCNsigmaPi; + bool is_ka_included_TPC = fHadronCut.cfg_min_TPCNsigmaKa < track.tpcNSigmaKa() && track.tpcNSigmaKa() < fHadronCut.cfg_max_TPCNsigmaKa; + bool is_pr_included_TPC = fHadronCut.cfg_min_TPCNsigmaPr < track.tpcNSigmaPr() && track.tpcNSigmaPr() < fHadronCut.cfg_max_TPCNsigmaPr; + return is_pi_included_TPC || is_ka_included_TPC || is_pr_included_TPC; + } + + template + bool isSelectedV0Leg(TTrack const& track) + { + if constexpr (isMC) { + if (!track.has_mcParticle()) { + return false; + } + } + + if (!track.hasITS() || !track.hasTPC()) { + return false; + } + + if (track.itsChi2NCl() > fV0Cut.cfg_max_chi2its) { + return false; + } + + if (track.itsNCls() < fV0Cut.cfg_min_ncluster_its) { + return false; + } + + if (track.itsNClsInnerBarrel() < fV0Cut.cfg_min_ncluster_itsib) { + return false; + } + + if (track.tpcChi2NCl() > fV0Cut.cfg_max_chi2tpc) { + return false; + } + + if (track.tpcNClsFound() < fV0Cut.cfg_min_ncluster_tpc) { + return false; + } + + if (track.tpcNClsCrossedRows() < fV0Cut.cfg_min_ncrossedrows_tpc) { + return false; + } + + if (track.tpcCrossedRowsOverFindableCls() < fV0Cut.cfg_min_cr2findable_ratio_tpc) { + return false; + } + + if (track.tpcFractionSharedCls() > fV0Cut.cfg_max_frac_shared_clusters_tpc) { + return false; + } + + return true; + } + + template + bool isPion(TTrack const& track) + { + return fV0Cut.cfg_min_TPCNsigmaPi < track.tpcNSigmaPi() && track.tpcNSigmaPi() < fV0Cut.cfg_max_TPCNsigmaPi; + } + + template + bool isKaon(TTrack const& track) + { + return fV0Cut.cfg_min_TPCNsigmaKa < track.tpcNSigmaKa() && track.tpcNSigmaKa() < fV0Cut.cfg_max_TPCNsigmaKa; + } + + template + bool isProton(TTrack const& track) + { + return fV0Cut.cfg_min_TPCNsigmaPr < track.tpcNSigmaPr() && track.tpcNSigmaPr() < fV0Cut.cfg_max_TPCNsigmaPr; + } + + template + bool isK0S(TV0 const& v0) + { + return (fV0Cut.cfg_min_mass_k0s < v0.mK0Short() && v0.mK0Short() < fV0Cut.cfg_max_mass_k0s) && (v0.mLambda() < fV0Cut.cfg_min_mass_lambda_veto || fV0Cut.cfg_max_mass_lambda_veto < v0.mLambda()) && (v0.mAntiLambda() < fV0Cut.cfg_min_mass_lambda_veto || fV0Cut.cfg_max_mass_lambda_veto < v0.mAntiLambda()); + } + + template + bool isLambda(TV0 const& v0) + { + return (fV0Cut.cfg_min_mass_lambda < v0.mLambda() && v0.mLambda() < fV0Cut.cfg_max_mass_lambda) && (v0.mK0Short() < fV0Cut.cfg_min_mass_k0s_veto || fV0Cut.cfg_max_mass_k0s_veto < v0.mK0Short()); + } + + template + bool isAntiLambda(TV0 const& v0) + { + return (fV0Cut.cfg_min_mass_lambda < v0.mAntiLambda() && v0.mAntiLambda() < fV0Cut.cfg_max_mass_lambda) && (v0.mK0Short() < fV0Cut.cfg_min_mass_k0s_veto || fV0Cut.cfg_max_mass_k0s_veto < v0.mK0Short()); + } + + template + bool isXi(TCascade const& cascade) + { + return (fCascadeCut.cfg_min_mass_Xi < cascade.mXi() && cascade.mXi() < fCascadeCut.cfg_max_mass_Xi) && (cascade.mOmega() < fCascadeCut.cfg_min_mass_Omega_veto || fCascadeCut.cfg_max_mass_Omega_veto < cascade.mOmega()); + } + + template + bool isOmega(TCascade const& cascade) + { + return (fCascadeCut.cfg_min_mass_Omega < cascade.mOmega() && cascade.mOmega() < fCascadeCut.cfg_max_mass_Omega) && (cascade.mXi() < fCascadeCut.cfg_min_mass_Xi_veto || fCascadeCut.cfg_max_mass_Xi_veto < cascade.mXi()); + } + + template + void fillV0Histograms(TCollision const&, TV0 const& v0, THistoregistry& registry) + { + auto pos = v0.template posTrack_as(); + auto neg = v0.template negTrack_as(); + registry.fill(HIST("SCT/V0/hPt"), v0.pt()); + registry.fill(HIST("SCT/V0/hAP"), v0.alpha(), v0.qtarm()); + registry.fill(HIST("SCT/V0/hCosPA"), v0.v0cosPA()); + registry.fill(HIST("SCT/V0/hLxy"), v0.v0radius()); + registry.fill(HIST("SCT/V0/hDCA2Legs"), v0.dcaV0daughters()); + + if (isPion(pos) && isPion(neg)) { + if ((v0.mLambda() < fV0Cut.cfg_min_mass_lambda_veto || fV0Cut.cfg_max_mass_lambda_veto < v0.mLambda()) && (v0.mAntiLambda() < fV0Cut.cfg_min_mass_lambda_veto || fV0Cut.cfg_max_mass_lambda_veto < v0.mAntiLambda())) { + registry.fill(HIST("SCT/V0/hMassK0S"), v0.mK0Short()); + registry.fill(HIST("SCT/V0/hYPhi_K0S"), v0.phi(), v0.yK0Short()); + } + registry.fill(HIST("SCT/V0/hMassGamma_misid"), v0.mGamma()); + registry.fill(HIST("SCT/V0/hMassLambda_misid"), v0.mLambda()); + registry.fill(HIST("SCT/V0/hMassAntiLambda_misid"), v0.mAntiLambda()); + } + + if (isProton(pos) && isPion(neg)) { + if (v0.mK0Short() < fV0Cut.cfg_min_mass_k0s_veto || fV0Cut.cfg_max_mass_k0s_veto < v0.mK0Short()) { + registry.fill(HIST("SCT/V0/hMassLambda"), v0.mLambda()); + registry.fill(HIST("SCT/V0/hYPhi_Lambda"), v0.phi(), v0.yLambda()); + } + registry.fill(HIST("SCT/V0/hMassGamma_misid"), v0.mGamma()); + registry.fill(HIST("SCT/V0/hMassK0S_misid"), v0.mK0Short()); + } + + if (isProton(neg) && isPion(pos)) { + if (v0.mK0Short() < fV0Cut.cfg_min_mass_k0s_veto || fV0Cut.cfg_max_mass_k0s_veto < v0.mK0Short()) { + registry.fill(HIST("SCT/V0/hMassAntiLambda"), v0.mAntiLambda()); + registry.fill(HIST("SCT/V0/hYPhi_Lambda"), v0.phi(), v0.yLambda()); + } + registry.fill(HIST("SCT/V0/hMassGamma_misid"), v0.mGamma()); + registry.fill(HIST("SCT/V0/hMassK0S_misid"), v0.mK0Short()); + } + } + + template + void fillCascadeHistograms(TCollision const& collision, TCascade const& cascade, THistoregistry& registry) + { + auto pos = cascade.template posTrack_as(); + auto neg = cascade.template negTrack_as(); + auto bachelor = cascade.template bachelor_as(); + + registry.fill(HIST("SCT/Cascade/hPt"), cascade.pt()); + registry.fill(HIST("SCT/Cascade/hMassLambda"), cascade.mLambda()); + registry.fill(HIST("SCT/Cascade/hCosPA"), cascade.casccosPA(collision.posX(), collision.posY(), collision.posZ())); + registry.fill(HIST("SCT/Cascade/hDCA2Legs"), cascade.dcacascdaughters()); + registry.fill(HIST("SCT/Cascade/hV0CosPA"), cascade.v0cosPA(collision.posX(), collision.posY(), collision.posZ())); + registry.fill(HIST("SCT/Cascade/hV0DCA2Legs"), cascade.dcaV0daughters()); + + if (cascade.sign() < 0) { // Xi- or Omega- + if (isPion(bachelor) && isProton(pos) && isPion(neg)) { + if (cascade.mOmega() < fCascadeCut.cfg_min_mass_Omega_veto || fCascadeCut.cfg_max_mass_Omega_veto < cascade.mOmega()) { + registry.fill(HIST("SCT/Cascade/hMassXi"), cascade.mXi()); + registry.fill(HIST("SCT/Cascade/hYPhi_Xi"), cascade.phi(), cascade.yXi()); + } + registry.fill(HIST("SCT/Cascade/hMassOmega_misid"), cascade.mOmega()); + } + if (isKaon(bachelor) && isProton(pos) && isPion(neg)) { + if (cascade.mXi() < fCascadeCut.cfg_min_mass_Xi_veto || fCascadeCut.cfg_max_mass_Xi_veto < cascade.mXi()) { + registry.fill(HIST("SCT/Cascade/hMassOmega"), cascade.mOmega()); + registry.fill(HIST("SCT/Cascade/hYPhi_Omega"), cascade.phi(), cascade.yOmega()); + } + registry.fill(HIST("SCT/Cascade/hMassXi_misid"), cascade.mXi()); + } + } else { // Xi+ or Omega+ + if (isPion(bachelor) && isProton(neg) && isPion(pos)) { + if (cascade.mOmega() < fCascadeCut.cfg_min_mass_Omega_veto || fCascadeCut.cfg_max_mass_Omega_veto < cascade.mOmega()) { + registry.fill(HIST("SCT/Cascade/hMassXi"), cascade.mXi()); + registry.fill(HIST("SCT/Cascade/hYPhi_Xi"), cascade.phi(), cascade.yXi()); + } + registry.fill(HIST("SCT/Cascade/hMassOmega_misid"), cascade.mOmega()); + } + if (isKaon(bachelor) && isProton(neg) && isPion(pos)) { + if (cascade.mXi() < fCascadeCut.cfg_min_mass_Xi_veto || fCascadeCut.cfg_max_mass_Xi_veto < cascade.mXi()) { + registry.fill(HIST("SCT/Cascade/hMassOmega"), cascade.mOmega()); + registry.fill(HIST("SCT/Cascade/hYPhi_Omega"), cascade.phi(), cascade.yOmega()); + } + registry.fill(HIST("SCT/Cascade/hMassXi_misid"), cascade.mXi()); + } + } + } + + template + void processWithTTCA(TBCs const& bcs, TCollisions const& collisions, TTracks const& tracks, TV0s const& v0s, TCascades const& cascades, TTrackAssoc const& trackIndices, TMCCollisions const& mcCollisions, TMCParticles const& mcParticles, TProducts& products, THistoregistry& registry, TSliceCache& cache, TPresliceTrack const& perColTrack, TPresliceTrackAssoc const& trackIndicesPerCollision, TPresliceV0 const& perColV0, TPresliceCascade const& perColCasc) + { + initCCDB(bcs.begin()); + + calculateTOFNSigmaWithReassociation(collisions, bcs, tracks, trackIndices, cache, perColTrack, trackIndicesPerCollision); + + for (const auto& collision : collisions) { + auto bc = collision.template bc_as(); + initCCDB(bc); + + if constexpr (isMC) { + if (!collision.has_mcCollision()) { + continue; + } + } + + if (!collision.isSelected()) { + continue; + } + + if constexpr (isTriggerAnalysis) { + if (collision.swtaliastmp_raw() == 0) { + continue; + } + } + + mVtx.setPos({collision.posX(), collision.posY(), collision.posZ()}); + mVtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ()); + std::unordered_multimap multiMapTracksPerCollision; // collisionId -> trackIds + + auto trackIdsThisCollision = trackIndices.sliceBy(trackIndicesPerCollision, collision.globalIndex()); + + std::vector looseElectronIds; // for pion prefilter + if (fElectronPFCut.doPF) { + looseElectronIds.reserve(trackIdsThisCollision.size()); + } + + for (const auto& trackId : trackIdsThisCollision) { + auto track = trackId.template track_as(); + if (!checkTrack(collision, track) && !checkTrackPF(collision, track)) { + continue; + } + + fillMapMLPID(collision, track); // call this function for all track also for loose electron sample. + + if (checkTrack(collision, track) && isElectron(collision, track, fElectronCut)) { + multiMapTracksPerCollision.insert(std::make_pair(collision.globalIndex(), track.globalIndex())); + } + + if (fElectronPFCut.doPF && checkTrackPF(collision, track) && isElectron(collision, track, fElectronPFCut)) { + looseElectronIds.emplace_back(track.globalIndex()); + } + } // end of track loop + + int count_electrons = multiMapTracksPerCollision.count(collision.globalIndex()); + registry.fill(HIST("Track/hNe"), count_electrons); + if (count_electrons < fElectronCut.minNelectron) { + continue; + } + + std::vector hadronIds; + std::vector k0sIds; + std::vector lambdaIds; + std::vector antilambdaIds; + std::vector xiMinusIds; + std::vector xiPlusIds; + std::vector omegaMinusIds; + std::vector omegaPlusIds; + + if (fDoSCTwithTracks) { + hadronIds.reserve(trackIdsThisCollision.size()); + } + + if (fDoSCTwithV0s) { + auto v0s_per_collision = v0s.sliceBy(perColV0, collision.globalIndex()); + k0sIds.reserve(v0s_per_collision.size()); + lambdaIds.reserve(v0s_per_collision.size()); + antilambdaIds.reserve(v0s_per_collision.size()); + + for (const auto& v0 : v0s_per_collision) { + auto pos = v0.template posTrack_as(); + auto neg = v0.template negTrack_as(); + + if (pos.sign() * neg.sign() > 0) { + continue; + } + if (!isSelectedV0Leg(pos) || !isSelectedV0Leg(neg)) { + continue; + } + + if (v0.dcaV0daughters() > fV0Cut.cfg_max_dca2legs) { + continue; + } + + if (v0.v0radius() < fV0Cut.cfg_min_radius) { + continue; + } + + if (v0.v0cosPA() < fV0Cut.cfg_min_cospa) { + continue; + } + + if (std::sqrt(std::pow(v0.alpha() / fV0Cut.cfg_max_alpha_veto, 2) + std::pow(v0.qtarm() / fV0Cut.cfg_max_qt_veto, 2)) < 1.f) { // photon conversion rejection at small qT + continue; + } + + fillV0Histograms(collision, v0, registry); + + if (isK0S(v0) && isPion(pos) && isPion(neg)) { + k0sIds.emplace_back(v0.globalIndex()); + } + + if (isLambda(v0) && isProton(pos) && isPion(neg)) { + lambdaIds.emplace_back(v0.globalIndex()); + } else if (isAntiLambda(v0) && isProton(neg) && isPion(pos)) { + antilambdaIds.emplace_back(v0.globalIndex()); + } + + } // end of v0 loop + } + + if (fDoSCTwithCascades) { + auto cascades_per_collision = cascades.sliceBy(perColCasc, collision.globalIndex()); + xiMinusIds.reserve(cascades_per_collision.size()); + xiPlusIds.reserve(cascades_per_collision.size()); + omegaMinusIds.reserve(cascades_per_collision.size()); + omegaPlusIds.reserve(cascades_per_collision.size()); + + for (const auto& cascade : cascades_per_collision) { + auto pos = cascade.template posTrack_as(); + auto neg = cascade.template negTrack_as(); + auto bachelor = cascade.template bachelor_as(); + if (pos.sign() * neg.sign() > 0) { + continue; + } + if (cascade.mLambda() < fCascadeCut.cfg_min_mass_lambda || fCascadeCut.cfg_max_mass_lambda < cascade.mLambda()) { + continue; + } + + if (!isSelectedV0Leg(pos) || !isSelectedV0Leg(neg) || !isSelectedV0Leg(bachelor)) { + continue; + } + + if (cascade.casccosPA(collision.posX(), collision.posY(), collision.posZ()) < fCascadeCut.cfg_min_cospa) { + continue; + } + if (cascade.v0cosPA(collision.posX(), collision.posY(), collision.posZ()) < fCascadeCut.cfg_min_cospa_v0) { + continue; + } + + if (cascade.cascradius() > cascade.v0radius()) { + continue; + } + + if (cascade.dcaV0daughters() > fCascadeCut.cfg_max_dcadau_v0) { + continue; + } + if (cascade.v0radius() < fCascadeCut.cfg_min_rxy_v0) { + continue; + } + if (cascade.cascradius() < fCascadeCut.cfg_min_rxy) { + continue; + } + + if (cascade.dcacascdaughters() > fCascadeCut.cfg_max_dcadau) { + continue; + } + + if (std::fabs(cascade.dcav0topv(collision.posX(), collision.posY(), collision.posZ())) < fCascadeCut.cfg_min_dcaxy_v0) { + continue; + } + + fillCascadeHistograms(collision, cascade, registry); + + if (cascade.sign() < 0) { // Xi- or Omega- + if (isXi(cascade) && isPion(bachelor) && isProton(pos) && isPion(neg)) { + xiMinusIds.emplace_back(cascade.globalIndex()); + } + if (isOmega(cascade) && isKaon(bachelor) && isProton(pos) && isPion(neg)) { + omegaMinusIds.emplace_back(cascade.globalIndex()); + } + } else { // Xi+ or Omega+ + if (isXi(cascade) && isPion(bachelor) && isProton(neg) && isPion(pos)) { + xiPlusIds.emplace_back(cascade.globalIndex()); + } + if (isOmega(cascade) && isKaon(bachelor) && isProton(neg) && isPion(pos)) { + omegaPlusIds.emplace_back(cascade.globalIndex()); + } + } + } // end of cascade loop + } + + for (const auto& trackId : trackIdsThisCollision) { + auto track = trackId.template track_as(); + auto trackParCov = getTrackParCov(track); + o2::dataformats::DCA mDcaInfoCov; + mDcaInfoCov.set(999, 999, 999, 999, 999); + trackParCov.setPID(track.pidForTracking()); + bool isPropOK = o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov); + if (!isPropOK) { + continue; + } + float dcaXY = mDcaInfoCov.getY(); + float dcaZ = mDcaInfoCov.getZ(); + if (isSelectedHadron(collision, track, trackParCov, dcaXY, dcaZ)) { + float tpcSignal = track.tpcSignal(); + if constexpr (isMC) { + tpcSignal = track.mcTunedTPCSignal(); + } + registry.fill(HIST("SCT/Track/hs"), trackParCov.getPt(), trackParCov.getEta(), RecoDecay::constrainAngle(trackParCov.getPhi(), 0, 1U)); + registry.fill(HIST("SCT/Track/hDCA"), dcaXY, dcaZ); + registry.fill(HIST("SCT/Track/hTPCdEdx"), track.tpcInnerParam(), tpcSignal); + registry.fill(HIST("SCT/Track/hTOFbeta"), trackParCov.getP(), fMapTOFBetaReassociated[std::make_pair(collision.globalIndex(), track.globalIndex())]); + hadronIds.emplace_back(track.globalIndex()); + } + } // end of track loop + + auto range_electrons = multiMapTracksPerCollision.equal_range(collision.globalIndex()); + for (auto it = range_electrons.first; it != range_electrons.second; it++) { + auto electron = tracks.rawIteratorAt(it->second); + o2::dataformats::DCA mDcaInfoCov; + mDcaInfoCov.set(999, 999, 999, 999, 999); + auto trackParCov = getTrackParCov(electron); + trackParCov.setPID(o2::track::PID::Electron); + bool isPropOK = o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov); + if (!isPropOK) { + continue; + } + float dcaXY = mDcaInfoCov.getY(); + float dcaZ = mDcaInfoCov.getZ(); + fillElectronTable(collision, electron, trackParCov, dcaXY, dcaZ, products, registry); + ROOT::Math::PtEtaPhiMVector v1(trackParCov.getPt(), trackParCov.getEta(), RecoDecay::constrainAngle(trackParCov.getPhi(), 0, 1U), o2::constants::physics::MassElectron); // main electron + + // apply pion prefilter + uint8_t pfb = 0; + for (const auto& looseElectronId : looseElectronIds) { + auto looseElectron = tracks.rawIteratorAt(looseElectronId); + if (looseElectron.globalIndex() == electron.globalIndex()) { + continue; + } + if (looseElectron.sign() * electron.sign() > 0) { // reject LS, because pi0dalitz and photon conversion must be ULS. + continue; + } + + auto trackParCov2 = getTrackParCov(looseElectron); + trackParCov2.setPID(o2::track::PID::Electron); + bool isPropOK = o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov2, 2.f, matCorr, &mDcaInfoCov); + if (!isPropOK) { + continue; + } + + float tpcSignal = looseElectron.tpcSignal(); + if constexpr (isMC) { + tpcSignal = looseElectron.mcTunedTPCSignal(); + } + registry.fill(HIST("Prefilter/Track/hPt"), trackParCov2.getPt()); + registry.fill(HIST("Prefilter/Track/hEtaPhi"), RecoDecay::constrainAngle(trackParCov2.getPhi(), 0, 1U), trackParCov2.getEta()); + registry.fill(HIST("Prefilter/Track/hTPCdEdx"), looseElectron.tpcInnerParam(), tpcSignal); + registry.fill(HIST("Prefilter/Track/hTOFbeta"), trackParCov2.getP(), fMapTOFBetaReassociated[std::make_pair(collision.globalIndex(), looseElectron.globalIndex())]); + + ROOT::Math::PtEtaPhiMVector v2(trackParCov2.getPt(), trackParCov2.getEta(), RecoDecay::constrainAngle(trackParCov2.getPhi(), 0, 1U), o2::constants::physics::MassElectron); // loose electron + ROOT::Math::PtEtaPhiMVector v12 = v1 + v2; + float phiv = o2::aod::pwgem::dilepton::utils::pairutil::getPhivPair(v1.Px(), v1.Py(), v1.Pz(), v2.Px(), v2.Py(), v2.Pz(), electron.sign(), looseElectron.sign(), d_bz); + registry.fill(HIST("Prefilter/Pair/hMvsPhiV"), phiv, v12.M()); + + if (v12.M() < fElectronPFCut.slope * phiv + fElectronPFCut.intercept) { + pfb |= (uint8_t(1) << static_cast(o2::aod::pwgem::dilepton::utils::pairutil::DileptonPrefilterBit::kElFromPC)); + } + + for (int i = 0; i < static_cast(max_mee_vec.size()); i++) { + if (v12.M() < max_mee_vec.at(i)) { + pfb |= (uint8_t(1) << (static_cast(o2::aod::pwgem::dilepton::utils::pairutil::DileptonPrefilterBit::kElFromPi0_20MeV) + i)); + } + } // end of mee loop + + } // end of loose electron loop + products.electronPFTable(pfb); + + // perform SCT from here + std::vector bdtScorePrompt; + std::vector bdtScorePromptHc; + std::vector bdtScoreNonpromptHc; + std::vector bdtScoreHb; + std::vector hadronType; + + bdtScorePrompt.reserve(hadronIds.size() + k0sIds.size() + lambdaIds.size() + antilambdaIds.size()); + bdtScorePromptHc.reserve(hadronIds.size() + k0sIds.size() + lambdaIds.size() + antilambdaIds.size()); + bdtScoreNonpromptHc.reserve(hadronIds.size() + k0sIds.size() + lambdaIds.size() + antilambdaIds.size()); + bdtScoreHb.reserve(hadronIds.size() + k0sIds.size() + lambdaIds.size() + antilambdaIds.size()); + hadronType.reserve(hadronIds.size() + k0sIds.size() + lambdaIds.size() + antilambdaIds.size()); + + // eTrack pair + for (const auto& hadronId : hadronIds) { + auto hadron = tracks.rawIteratorAt(hadronId); + if (hadron.globalIndex() == electron.globalIndex()) { + continue; + } + + auto hadronParCov = getTrackParCov(hadron); + o2::dataformats::DCA mDcaInfoCov; + mDcaInfoCov.set(999, 999, 999, 999, 999); + hadronParCov.setPID(hadron.pidForTracking()); + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, hadronParCov, 2.f, matCorr, &mDcaInfoCov); + + auto eTpair = o2::aod::pwgem::dilepton::utils::makePairLeptonTrack(dfeT, collision, electron, hadron, o2::track::PID::Electron, hadron.pidForTracking()); + registry.fill(HIST("SCT/eT/hImpactParameter"), eTpair.impParXY, eTpair.impParZ); + registry.fill(HIST("SCT/eT/hDecayLength"), eTpair.lxy, eTpair.lz); + registry.fill(HIST("SCT/eT/hCosPA"), eTpair.cospa); + registry.fill(HIST("SCT/eT/hDCA2legs"), eTpair.dca2legs); + registry.fill(HIST("SCT/eT/hMass"), eTpair.mass); + if (eTpair.isOK && fConfigDFeT.useML) { + o2::analysis::pwgem::dilepton::sct::candidate candidate; + fillCandidate(candidate, eTpair, trackParCov, mDcaInfoCov); + candidate.ptL = trackParCov.getPt(); + candidate.signLH = electron.sign() * hadron.sign(); + candidate.signedMassLH = electron.sign() * hadron.sign() * eTpair.mass; + candidate.tpcNSigmaKa = hadron.tpcNSigmaKa(); + + std::vector inputFeatures = mlResponseSCTeT.getInputFeatures(candidate); + float binningFeature = mlResponseSCTeT.getBinningFeature(candidate); + + int pbin = lower_bound(fConfigDFeT.binsMl.value.begin(), fConfigDFeT.binsMl.value.end(), binningFeature) - fConfigDFeT.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(fConfigDFeT.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(fConfigDFeT.binsMl.value.size()) - 2; + } + + float probaPrompt = mlResponseSCTeT.getModelOutput(inputFeatures, pbin)[0]; + float probaPromptHc = mlResponseSCTeT.getModelOutput(inputFeatures, pbin)[1]; + float probaNonpromptHc = mlResponseSCTeT.getModelOutput(inputFeatures, pbin)[2]; + float probaHb = mlResponseSCTeT.getModelOutput(inputFeatures, pbin)[3]; + + bdtScorePrompt.emplace_back(probaPrompt); + bdtScorePromptHc.emplace_back(probaPromptHc); + bdtScoreNonpromptHc.emplace_back(probaNonpromptHc); + bdtScoreHb.emplace_back(probaHb); + hadronType.emplace_back(0); + } + } // end of charged track loop + + // eK0S pair + for (const auto& k0sId : k0sIds) { + auto v0 = v0s.rawIteratorAt(k0sId); + if (v0.posTrackId() == electron.globalIndex() || v0.negTrackId() == electron.globalIndex()) { + continue; + } + const std::array vertexV0 = {v0.x(), v0.y(), v0.z()}; + const std::array momV0 = {v0.px(), v0.py(), v0.pz()}; + std::array covV0 = {0.}; + constexpr int MomInd[6] = {9, 13, 14, 18, 19, 20}; // cov matrix elements for momentum component + for (int i = 0; i < 6; i++) { + covV0[MomInd[i]] = v0.momentumCovMat()[i]; + covV0[i] = v0.positionCovMat()[i]; + } + auto v0ParCov = o2::track::TrackParCov(vertexV0, momV0, covV0, 0, true); + v0ParCov.setAbsCharge(0); + v0ParCov.setPID(o2::track::PID::K0); + o2::dataformats::DCA impactParameterV0; + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, v0ParCov, 2.f, matCorr, &impactParameterV0); // v0ParCov is TrackParCov object + + auto eV0pair = o2::aod::pwgem::dilepton::utils::makePairLeptonV0(dfeV0, collision, electron, v0, o2::track::PID::Electron, o2::track::PID::K0); + registry.fill(HIST("SCT/eV0/hImpactParameter"), eV0pair.impParXY, eV0pair.impParZ); + registry.fill(HIST("SCT/eV0/hDecayLength"), eV0pair.lxy, eV0pair.lz); + registry.fill(HIST("SCT/eV0/hCosPA"), eV0pair.cospa); + registry.fill(HIST("SCT/eV0/hDCA2legs"), eV0pair.dca2legs); + registry.fill(HIST("SCT/eV0/hMass"), eV0pair.mass); + if (eV0pair.isOK && fConfigDFeV0.useML) { + o2::analysis::pwgem::dilepton::sct::candidate candidate; + fillCandidate(candidate, eV0pair, v0ParCov, impactParameterV0); + candidate.ptL = trackParCov.getPt(); + + std::vector inputFeatures = mlResponseSCTeV0.getInputFeatures(candidate); + float binningFeature = mlResponseSCTeV0.getBinningFeature(candidate); + + int pbin = lower_bound(fConfigDFeV0.binsMl.value.begin(), fConfigDFeV0.binsMl.value.end(), binningFeature) - fConfigDFeV0.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(fConfigDFeV0.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(fConfigDFeV0.binsMl.value.size()) - 2; + } + + float probaPrompt = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[0]; + float probaPromptHc = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[1]; + float probaNonpromptHc = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[2]; + float probaHb = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[3]; + bdtScorePrompt.emplace_back(probaPrompt); + bdtScorePromptHc.emplace_back(probaPromptHc); + bdtScoreNonpromptHc.emplace_back(probaNonpromptHc); + bdtScoreHb.emplace_back(probaHb); + hadronType.emplace_back(1); + } + } // end of k0s loop + + if (electron.sign() > 0) { // positron + // eL pair // sign is restricted in baryon decay: Lc+ -> e+ nu_e Lambda + for (const auto& lambdaId : lambdaIds) { + auto v0 = v0s.rawIteratorAt(lambdaId); + if (v0.posTrackId() == electron.globalIndex() || v0.negTrackId() == electron.globalIndex()) { + continue; + } + const std::array vertexV0 = {v0.x(), v0.y(), v0.z()}; + const std::array momV0 = {v0.px(), v0.py(), v0.pz()}; + std::array covV0 = {0.}; + constexpr int MomInd[6] = {9, 13, 14, 18, 19, 20}; // cov matrix elements for momentum component + for (int i = 0; i < 6; i++) { + covV0[MomInd[i]] = v0.momentumCovMat()[i]; + covV0[i] = v0.positionCovMat()[i]; + } + auto v0ParCov = o2::track::TrackParCov(vertexV0, momV0, covV0, 0, true); + v0ParCov.setAbsCharge(0); + v0ParCov.setPID(o2::track::PID::Lambda); + o2::dataformats::DCA impactParameterV0; + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, v0ParCov, 2.f, matCorr, &impactParameterV0); // v0ParCov is TrackParCov object + + auto eV0pair = o2::aod::pwgem::dilepton::utils::makePairLeptonV0(dfeV0, collision, electron, v0, o2::track::PID::Electron, o2::track::PID::Lambda); + registry.fill(HIST("SCT/eV0/hImpactParameter"), eV0pair.impParXY, eV0pair.impParZ); + registry.fill(HIST("SCT/eV0/hDecayLength"), eV0pair.lxy, eV0pair.lz); + registry.fill(HIST("SCT/eV0/hCosPA"), eV0pair.cospa); + registry.fill(HIST("SCT/eV0/hDCA2legs"), eV0pair.dca2legs); + registry.fill(HIST("SCT/eV0/hMass"), eV0pair.mass); + if (eV0pair.isOK && fConfigDFeV0.useML) { + o2::analysis::pwgem::dilepton::sct::candidate candidate; + fillCandidate(candidate, eV0pair, v0ParCov, impactParameterV0); + candidate.ptL = trackParCov.getPt(); + + std::vector inputFeatures = mlResponseSCTeV0.getInputFeatures(candidate); + float binningFeature = mlResponseSCTeV0.getBinningFeature(candidate); + + int pbin = lower_bound(fConfigDFeV0.binsMl.value.begin(), fConfigDFeV0.binsMl.value.end(), binningFeature) - fConfigDFeV0.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(fConfigDFeV0.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(fConfigDFeV0.binsMl.value.size()) - 2; + } + + float probaPrompt = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[0]; + float probaPromptHc = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[1]; + float probaNonpromptHc = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[2]; + float probaHb = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[3]; + bdtScorePrompt.emplace_back(probaPrompt); + bdtScorePromptHc.emplace_back(probaPromptHc); + bdtScoreNonpromptHc.emplace_back(probaNonpromptHc); + bdtScoreHb.emplace_back(probaHb); + hadronType.emplace_back(2); + } + } // end of lambda loop + + // eXi pair // sign is restricted in baryon decay: Xic0 -> e+ nu_e Xi- + for (const auto& xiId : xiMinusIds) { + auto cascade = cascades.rawIteratorAt(xiId); + if (cascade.posTrackId() == electron.globalIndex() || cascade.negTrackId() == electron.globalIndex() || cascade.bachelorId() == electron.globalIndex()) { + continue; + } + + const std::array vertexCasc = {cascade.x(), cascade.y(), cascade.z()}; + const std::array momCasc = {cascade.px(), cascade.py(), cascade.pz()}; + std::array covCasc = {0.}; + constexpr int MomInd[6] = {9, 13, 14, 18, 19, 20}; // cov matrix elements for momentum component + for (int i = 0; i < 6; i++) { + covCasc[MomInd[i]] = cascade.momentumCovMat()[i]; + covCasc[i] = cascade.positionCovMat()[i]; + } + auto cascadeParCov = o2::track::TrackParCov(vertexCasc, momCasc, covCasc, cascade.sign(), true); + cascadeParCov.setAbsCharge(1); + cascadeParCov.setPID(o2::track::PID::XiMinus); + o2::dataformats::DCA impactParameterCasc; + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, cascadeParCov, 2.f, matCorr, &impactParameterCasc); // cascadeParCov is TrackParCov object + + auto eCpair = o2::aod::pwgem::dilepton::utils::makePairLeptonCascade(dfeC, collision, electron, cascade, o2::track::PID::Electron, o2::track::PID::XiMinus); + registry.fill(HIST("SCT/eC/hImpactParameter"), eCpair.impParXY, eCpair.impParZ); + registry.fill(HIST("SCT/eC/hDecayLength"), eCpair.lxy, eCpair.lz); + registry.fill(HIST("SCT/eC/hCosPA"), eCpair.cospa); + registry.fill(HIST("SCT/eC/hDCA2legs"), eCpair.dca2legs); + registry.fill(HIST("SCT/eC/hMass"), eCpair.mass); + if (eCpair.isOK && fConfigDFeC.useML) { + o2::analysis::pwgem::dilepton::sct::candidate candidate; + fillCandidate(candidate, eCpair, cascadeParCov, impactParameterCasc); + candidate.ptL = trackParCov.getPt(); + + std::vector inputFeatures = mlResponseSCTeC.getInputFeatures(candidate); + float binningFeature = mlResponseSCTeC.getBinningFeature(candidate); + + int pbin = lower_bound(fConfigDFeC.binsMl.value.begin(), fConfigDFeC.binsMl.value.end(), binningFeature) - fConfigDFeC.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(fConfigDFeC.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(fConfigDFeC.binsMl.value.size()) - 2; + } + + float probaPrompt = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[0]; + float probaPromptHc = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[1]; + float probaNonpromptHc = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[2]; + float probaHb = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[3]; + bdtScorePrompt.emplace_back(probaPrompt); + bdtScorePromptHc.emplace_back(probaPromptHc); + bdtScoreNonpromptHc.emplace_back(probaNonpromptHc); + bdtScoreHb.emplace_back(probaHb); + hadronType.emplace_back(4); + } + } // end of Xi- loop + + // eOmega pair // sign is restricted in baryon decay: Omegac0 -> e+ nu_e Omega- + for (const auto& omegaId : omegaMinusIds) { + auto cascade = cascades.rawIteratorAt(omegaId); + if (cascade.posTrackId() == electron.globalIndex() || cascade.negTrackId() == electron.globalIndex() || cascade.bachelorId() == electron.globalIndex()) { + continue; + } + + const std::array vertexCasc = {cascade.x(), cascade.y(), cascade.z()}; + const std::array momCasc = {cascade.px(), cascade.py(), cascade.pz()}; + std::array covCasc = {0.}; + constexpr int MomInd[6] = {9, 13, 14, 18, 19, 20}; // cov matrix elements for momentum component + for (int i = 0; i < 6; i++) { + covCasc[MomInd[i]] = cascade.momentumCovMat()[i]; + covCasc[i] = cascade.positionCovMat()[i]; + } + auto cascadeParCov = o2::track::TrackParCov(vertexCasc, momCasc, covCasc, cascade.sign(), true); + cascadeParCov.setAbsCharge(1); + cascadeParCov.setPID(o2::track::PID::OmegaMinus); + o2::dataformats::DCA impactParameterCasc; + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, cascadeParCov, 2.f, matCorr, &impactParameterCasc); // cascadeParCov is TrackParCov object + + auto eCpair = o2::aod::pwgem::dilepton::utils::makePairLeptonCascade(dfeC, collision, electron, cascade, o2::track::PID::Electron, o2::track::PID::OmegaMinus); + registry.fill(HIST("SCT/eC/hImpactParameter"), eCpair.impParXY, eCpair.impParZ); + registry.fill(HIST("SCT/eC/hDecayLength"), eCpair.lxy, eCpair.lz); + registry.fill(HIST("SCT/eC/hCosPA"), eCpair.cospa); + registry.fill(HIST("SCT/eC/hDCA2legs"), eCpair.dca2legs); + registry.fill(HIST("SCT/eC/hMass"), eCpair.mass); + if (eCpair.isOK && fConfigDFeC.useML) { + o2::analysis::pwgem::dilepton::sct::candidate candidate; + fillCandidate(candidate, eCpair, cascadeParCov, impactParameterCasc); + candidate.ptL = trackParCov.getPt(); + + std::vector inputFeatures = mlResponseSCTeC.getInputFeatures(candidate); + float binningFeature = mlResponseSCTeC.getBinningFeature(candidate); + + int pbin = lower_bound(fConfigDFeC.binsMl.value.begin(), fConfigDFeC.binsMl.value.end(), binningFeature) - fConfigDFeC.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(fConfigDFeC.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(fConfigDFeC.binsMl.value.size()) - 2; + } + + float probaPrompt = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[0]; + float probaPromptHc = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[1]; + float probaNonpromptHc = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[2]; + float probaHb = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[3]; + bdtScorePrompt.emplace_back(probaPrompt); + bdtScorePromptHc.emplace_back(probaPromptHc); + bdtScoreNonpromptHc.emplace_back(probaNonpromptHc); + bdtScoreHb.emplace_back(probaHb); + hadronType.emplace_back(6); + } + } // end of Omega- loop + + } else { // electron + // eL pair // sign is restricted in baryon decay: Lc- -> e- anti_nu_e antiLambda + for (const auto& antilambdaId : antilambdaIds) { + auto v0 = v0s.rawIteratorAt(antilambdaId); + if (v0.posTrackId() == electron.globalIndex() || v0.negTrackId() == electron.globalIndex()) { + continue; + } + const std::array vertexV0 = {v0.x(), v0.y(), v0.z()}; + const std::array momV0 = {v0.px(), v0.py(), v0.pz()}; + std::array covV0 = {0.}; + constexpr int MomInd[6] = {9, 13, 14, 18, 19, 20}; // cov matrix elements for momentum component + for (int i = 0; i < 6; i++) { + covV0[MomInd[i]] = v0.momentumCovMat()[i]; + covV0[i] = v0.positionCovMat()[i]; + } + auto v0ParCov = o2::track::TrackParCov(vertexV0, momV0, covV0, 0, true); + v0ParCov.setAbsCharge(0); + v0ParCov.setPID(o2::track::PID::Lambda); + o2::dataformats::DCA impactParameterV0; + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, v0ParCov, 2.f, matCorr, &impactParameterV0); // v0ParCov is TrackParCov object + + auto eV0pair = o2::aod::pwgem::dilepton::utils::makePairLeptonV0(dfeV0, collision, electron, v0, o2::track::PID::Electron, o2::track::PID::Lambda); + registry.fill(HIST("SCT/eV0/hImpactParameter"), eV0pair.impParXY, eV0pair.impParZ); + registry.fill(HIST("SCT/eV0/hDecayLength"), eV0pair.lxy, eV0pair.lz); + registry.fill(HIST("SCT/eV0/hCosPA"), eV0pair.cospa); + registry.fill(HIST("SCT/eV0/hDCA2legs"), eV0pair.dca2legs); + registry.fill(HIST("SCT/eV0/hMass"), eV0pair.mass); + if (eV0pair.isOK && fConfigDFeV0.useML) { + o2::analysis::pwgem::dilepton::sct::candidate candidate; + fillCandidate(candidate, eV0pair, v0ParCov, impactParameterV0); + candidate.ptL = trackParCov.getPt(); + + std::vector inputFeatures = mlResponseSCTeV0.getInputFeatures(candidate); + float binningFeature = mlResponseSCTeV0.getBinningFeature(candidate); + + int pbin = lower_bound(fConfigDFeV0.binsMl.value.begin(), fConfigDFeV0.binsMl.value.end(), binningFeature) - fConfigDFeV0.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(fConfigDFeV0.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(fConfigDFeV0.binsMl.value.size()) - 2; + } + + float probaPrompt = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[0]; + float probaPromptHc = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[1]; + float probaNonpromptHc = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[2]; + float probaHb = mlResponseSCTeV0.getModelOutput(inputFeatures, pbin)[3]; + bdtScorePrompt.emplace_back(probaPrompt); + bdtScorePromptHc.emplace_back(probaPromptHc); + bdtScoreNonpromptHc.emplace_back(probaNonpromptHc); + bdtScoreHb.emplace_back(probaHb); + hadronType.emplace_back(3); + } + } // end of antilambda loop + + // eXi pair // sign is restricted in baryon decay: Xic0bar -> e- anti_nu_e Xi+ + for (const auto& xiId : xiPlusIds) { + auto cascade = cascades.rawIteratorAt(xiId); + if (cascade.posTrackId() == electron.globalIndex() || cascade.negTrackId() == electron.globalIndex() || cascade.bachelorId() == electron.globalIndex()) { + continue; + } + + const std::array vertexCasc = {cascade.x(), cascade.y(), cascade.z()}; + const std::array momCasc = {cascade.px(), cascade.py(), cascade.pz()}; + std::array covCasc = {0.}; + constexpr int MomInd[6] = {9, 13, 14, 18, 19, 20}; // cov matrix elements for momentum component + for (int i = 0; i < 6; i++) { + covCasc[MomInd[i]] = cascade.momentumCovMat()[i]; + covCasc[i] = cascade.positionCovMat()[i]; + } + auto cascadeParCov = o2::track::TrackParCov(vertexCasc, momCasc, covCasc, cascade.sign(), true); + cascadeParCov.setAbsCharge(1); + cascadeParCov.setPID(o2::track::PID::XiMinus); + o2::dataformats::DCA impactParameterCasc; + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, cascadeParCov, 2.f, matCorr, &impactParameterCasc); // cascadeParCov is TrackParCov object + + auto eCpair = o2::aod::pwgem::dilepton::utils::makePairLeptonCascade(dfeC, collision, electron, cascade, o2::track::PID::Electron, o2::track::PID::XiMinus); + registry.fill(HIST("SCT/eC/hImpactParameter"), eCpair.impParXY, eCpair.impParZ); + registry.fill(HIST("SCT/eC/hDecayLength"), eCpair.lxy, eCpair.lz); + registry.fill(HIST("SCT/eC/hCosPA"), eCpair.cospa); + registry.fill(HIST("SCT/eC/hDCA2legs"), eCpair.dca2legs); + registry.fill(HIST("SCT/eC/hMass"), eCpair.mass); + if (eCpair.isOK && fConfigDFeC.useML) { + o2::analysis::pwgem::dilepton::sct::candidate candidate; + fillCandidate(candidate, eCpair, cascadeParCov, impactParameterCasc); + candidate.ptL = trackParCov.getPt(); + + std::vector inputFeatures = mlResponseSCTeC.getInputFeatures(candidate); + float binningFeature = mlResponseSCTeC.getBinningFeature(candidate); + + int pbin = lower_bound(fConfigDFeC.binsMl.value.begin(), fConfigDFeC.binsMl.value.end(), binningFeature) - fConfigDFeC.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(fConfigDFeC.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(fConfigDFeC.binsMl.value.size()) - 2; + } + + float probaPrompt = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[0]; + float probaPromptHc = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[1]; + float probaNonpromptHc = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[2]; + float probaHb = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[3]; + bdtScorePrompt.emplace_back(probaPrompt); + bdtScorePromptHc.emplace_back(probaPromptHc); + bdtScoreNonpromptHc.emplace_back(probaNonpromptHc); + bdtScoreHb.emplace_back(probaHb); + hadronType.emplace_back(5); + } + } // end of Xi- loop + + // eOmega pair // sign is restricted in baryon decay: Omegac0bar -> e- anti_nu_e Omega+ + for (const auto& omegaId : omegaPlusIds) { + auto cascade = cascades.rawIteratorAt(omegaId); + if (cascade.posTrackId() == electron.globalIndex() || cascade.negTrackId() == electron.globalIndex() || cascade.bachelorId() == electron.globalIndex()) { + continue; + } + + const std::array vertexCasc = {cascade.x(), cascade.y(), cascade.z()}; + const std::array momCasc = {cascade.px(), cascade.py(), cascade.pz()}; + std::array covCasc = {0.}; + constexpr int MomInd[6] = {9, 13, 14, 18, 19, 20}; // cov matrix elements for momentum component + for (int i = 0; i < 6; i++) { + covCasc[MomInd[i]] = cascade.momentumCovMat()[i]; + covCasc[i] = cascade.positionCovMat()[i]; + } + auto cascadeParCov = o2::track::TrackParCov(vertexCasc, momCasc, covCasc, cascade.sign(), true); + cascadeParCov.setAbsCharge(1); + cascadeParCov.setPID(o2::track::PID::OmegaMinus); + o2::dataformats::DCA impactParameterCasc; + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, cascadeParCov, 2.f, matCorr, &impactParameterCasc); // cascadeParCov is TrackParCov object + + auto eCpair = o2::aod::pwgem::dilepton::utils::makePairLeptonCascade(dfeC, collision, electron, cascade, o2::track::PID::Electron, o2::track::PID::OmegaMinus); + registry.fill(HIST("SCT/eC/hImpactParameter"), eCpair.impParXY, eCpair.impParZ); + registry.fill(HIST("SCT/eC/hDecayLength"), eCpair.lxy, eCpair.lz); + registry.fill(HIST("SCT/eC/hCosPA"), eCpair.cospa); + registry.fill(HIST("SCT/eC/hDCA2legs"), eCpair.dca2legs); + registry.fill(HIST("SCT/eC/hMass"), eCpair.mass); + if (eCpair.isOK && fConfigDFeC.useML) { + o2::analysis::pwgem::dilepton::sct::candidate candidate; + fillCandidate(candidate, eCpair, cascadeParCov, impactParameterCasc); + candidate.ptL = trackParCov.getPt(); + + std::vector inputFeatures = mlResponseSCTeC.getInputFeatures(candidate); + float binningFeature = mlResponseSCTeC.getBinningFeature(candidate); + + int pbin = lower_bound(fConfigDFeC.binsMl.value.begin(), fConfigDFeC.binsMl.value.end(), binningFeature) - fConfigDFeC.binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(fConfigDFeC.binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(fConfigDFeC.binsMl.value.size()) - 2; + } + + float probaPrompt = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[0]; + float probaPromptHc = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[1]; + float probaNonpromptHc = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[2]; + float probaHb = mlResponseSCTeC.getModelOutput(inputFeatures, pbin)[3]; + bdtScorePrompt.emplace_back(probaPrompt); + bdtScorePromptHc.emplace_back(probaPromptHc); + bdtScoreNonpromptHc.emplace_back(probaNonpromptHc); + bdtScoreHb.emplace_back(probaHb); + hadronType.emplace_back(7); + } + } // end of Omega- loop + } + + products.sctTable(bdtScorePrompt, bdtScorePromptHc, bdtScoreNonpromptHc, bdtScoreHb, hadronType); + + bdtScorePrompt.clear(); + bdtScorePromptHc.clear(); + bdtScoreNonpromptHc.clear(); + bdtScoreHb.clear(); + hadronType.clear(); + bdtScorePrompt.shrink_to_fit(); + bdtScorePromptHc.shrink_to_fit(); + bdtScoreNonpromptHc.shrink_to_fit(); + bdtScoreHb.shrink_to_fit(); + hadronType.shrink_to_fit(); + } // end of main electron loop + + hadronIds.clear(); + hadronIds.shrink_to_fit(); + + k0sIds.clear(); + k0sIds.shrink_to_fit(); + + lambdaIds.clear(); + lambdaIds.shrink_to_fit(); + antilambdaIds.clear(); + antilambdaIds.shrink_to_fit(); + + xiMinusIds.clear(); + xiMinusIds.shrink_to_fit(); + xiPlusIds.clear(); + xiPlusIds.shrink_to_fit(); + + omegaMinusIds.clear(); + omegaMinusIds.shrink_to_fit(); + omegaPlusIds.clear(); + omegaPlusIds.shrink_to_fit(); + + looseElectronIds.clear(); + looseElectronIds.shrink_to_fit(); + multiMapTracksPerCollision.clear(); + } // end of collision loop + clear(); + } + + template + void processWithoutTTCA(TBCs const& bcs, TCollisions const& collisions, TTracks const& tracks, TV0s const& v0s, TCascades const& cascades, TMCCollisions const& mcCollisions, TMCParticles const& mcParticles, TProducts& products, THistoregistry& registry) + { + LOGF(info, "processWithoutTTCA is not supported. Bye."); + return; + + initCCDB(bcs.begin()); + // calculateTOFNSigmaWithReassociation(collisions, bcs, tracks, nullptr); + + for (const auto& collision : collisions) { + auto bc = collision.template bc_as(); + initCCDB(bc); + + if constexpr (isMC) { + if (!collision.has_mcCollision()) { + continue; + } + } + + if (!collision.isSelected()) { + continue; + } + + if constexpr (isTriggerAnalysis) { + if (collision.swtaliastmp_raw() == 0) { + continue; + } + } + + } // end of collision loop + + clear(); + } + + void clear() + { + // this should be called at the end of DF. + fMapProbaEl.clear(); + fMapTOFNsigmaElReassociated.clear(); + fMapTOFNsigmaPiReassociated.clear(); + fMapTOFNsigmaKaReassociated.clear(); + fMapTOFNsigmaPrReassociated.clear(); + fMapTOFBetaReassociated.clear(); + } + + void doSCTwithTracks(const bool flag) { fDoSCTwithTracks = flag; } + void doSCTwithV0s(const bool flag) { fDoSCTwithV0s = flag; } + void doSCTwithCascades(const bool flag) { fDoSCTwithCascades = flag; } + + float dca3DinSigmaOTF(const float dcaXY, const float dcaZ, const float cYY, const float cZZ, const float cZY) + { + float det = cYY * cZZ - cZY * cZY; // determinant + if (det < 0) { + return 999.f; + } else { + return std::sqrt(std::fabs((dcaXY * dcaXY * cZZ + dcaZ * dcaZ * cYY - 2. * dcaXY * dcaZ * cZY) / det / 2.)); // dca 3d in sigma + } + } + + template + void fillCandidate(TCandidate& candidate, TPair const& pair, TTrackParCov const& trackParCov, TDCA const& dcaInfo) + { + candidate.ptH = trackParCov.getPt(); + candidate.tpcNSigmaKa = 0; + candidate.impParXYH = dcaInfo.getY(); + candidate.impParZH = dcaInfo.getZ(); + candidate.impPar3DH = std::sqrt(std::pow(candidate.impParXYH, 2) + std::pow(candidate.impParZH, 2)); + candidate.impParXYHinSigma = candidate.impParXYH / std::sqrt(trackParCov.getSigmaY2()); + candidate.impParZHinSigma = candidate.impParZH / std::sqrt(trackParCov.getSigmaZ2()); + candidate.impPar3DHinSigma = dca3DinSigmaOTF(candidate.impParXYH, candidate.impParZH, trackParCov.getSigmaY2(), trackParCov.getSigmaZ2(), trackParCov.getSigmaZY()); + candidate.signLH = 0; + candidate.dcaLH = pair.dca2legs; + candidate.massLH = pair.mass; + candidate.ptLH = pair.pt; + candidate.signedMassLH = pair.mass; + candidate.cpa = pair.cospa; + candidate.cpaXY = pair.cospaXY; + candidate.impParXY = pair.impParXY; + candidate.impParZ = pair.impParZ; + candidate.impPar3D = std::sqrt(std::pow(candidate.impParXY, 2) + std::pow(candidate.impParZ, 2)); + candidate.impParXYinSigma = candidate.impParXY / std::sqrt(pair.impParCYY); + candidate.impParZinSigma = candidate.impParZ / std::sqrt(pair.impParCZZ); + candidate.impPar3DinSigma = dca3DinSigmaOTF(candidate.impParXY, candidate.impParZ, pair.impParCYY, pair.impParCZY, pair.impParCZZ); + candidate.decayLengthXY = pair.lxy; + candidate.decayLengthZ = pair.lz; + candidate.decayLength3D = pair.lxyz; + candidate.decayLengthXYinSigma = pair.lxy / pair.lxyErr; + candidate.decayLengthZinSigma = pair.lz / pair.lzErr; + candidate.decayLength3DinSigma = pair.lxyz / pair.lxyzErr; + } + + protected: + electronCut fElectronCut; + electronPFCut fElectronPFCut; + hadronCut fHadronCut; + v0Cut fV0Cut; + cascadeCut fCascadeCut; + cfgDFeT fConfigDFeT; + cfgDFeV0 fConfigDFeV0; + cfgDFeC fConfigDFeC; + + std::map, float> fMapProbaEl; // map pair(collisionId, trackId) -> probaEl + std::map, float> fMapTOFNsigmaElReassociated; // map pair(collisionId, trackId) -> tof n sigma el + std::map, float> fMapTOFNsigmaPiReassociated; // map pair(collisionId, trackId) -> tof n sigma pi + std::map, float> fMapTOFNsigmaKaReassociated; // map pair(collisionId, trackId) -> tof n sigma ka + std::map, float> fMapTOFNsigmaPrReassociated; // map pair(collisionId, trackId) -> tof n sigma pr + std::map, float> fMapTOFBetaReassociated; // map pair(collisionId, trackId) -> tof beta + + int mRunNumber{0}; + float d_bz{0}; + o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrLUT; + o2::dataformats::VertexBase mVtx; + o2::framework::Service mTOFResponse; + o2::analysis::MlResponsePID mlResponsePID; + o2::analysis::MlResponseSCT mlResponseSCTeT; + o2::analysis::MlResponseSCT mlResponseSCTeV0; + o2::analysis::MlResponseSCT mlResponseSCTeC; + + bool fDoSCTwithTracks{false}; + bool fDoSCTwithV0s{false}; + bool fDoSCTwithCascades{false}; + const std::vector max_mee_vec{0.02, 0.04, 0.06, 0.08, 0.10, 0.12, 0.14}; + o2::vertexing::DCAFitterN<2> dfeT; + o2::vertexing::DCAFitterN<2> dfeV0; + o2::vertexing::DCAFitterN<2> dfeC; + o2::ccdb::CcdbApi ccdbApi; + +}; // end ElectronModule + +} // namespace pwgem::dilepton::utils +} // namespace o2 + +#endif // PWGEM_DILEPTON_UTILS_ELECTRONMODULE_H_ diff --git a/PWGEM/Dilepton/Utils/MlResponsePID.h b/PWGEM/Dilepton/Utils/MlResponsePID.h new file mode 100644 index 00000000000..63e85a6b063 --- /dev/null +++ b/PWGEM/Dilepton/Utils/MlResponsePID.h @@ -0,0 +1,147 @@ +// 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 MlResponsePID.h +/// \brief Class to compute the ML response for fwdtracks +/// \author Daiki Sekihata + +#ifndef PWGEM_DILEPTON_UTILS_MLRESPONSEPID_H_ +#define PWGEM_DILEPTON_UTILS_MLRESPONSEPID_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_TRACK(FEATURE) \ + { \ + #FEATURE, static_cast(InputFeaturesPID::FEATURE)} + +// Check if the index of mCachedIndices (index associated to a FEATURE) +// matches the entry in EnumInputFeatures associated to this FEATURE +// if so, the inputFeatures vector is filled with the FEATURE's value +// by calling the corresponding GETTER=FEATURE from track +#define CHECK_AND_FILL_TRACK(GETTER) \ + case static_cast(InputFeaturesPID::GETTER): { \ + inputFeature = track.GETTER; \ + break; \ + } + +namespace o2::analysis +{ +// possible input features for ML +enum class InputFeaturesPID : uint8_t { + tpcInnerParam, + tpcNClsFound, + tpcChi2NCl, + tpcNSigmaEl, + tofNSigmaEl, + meanClusterSizeITSobCosTgl, +}; + +namespace pwgem::dilepton::mlpid +{ +struct candidate { + float tpcInnerParam{0}; + int tpcNClsFound{0}; + float tpcChi2NCl{0}; + float tpcNSigmaEl{0}; + float tofNSigmaEl{0}; + float meanClusterSizeITSobCosTgl{0}; +}; +} // namespace pwgem::dilepton + +template +class MlResponsePID : public MlResponse +{ + public: + /// Default constructor + MlResponsePID() = default; + /// Default destructor + virtual ~MlResponsePID() = default; + + template + float return_feature(uint8_t idx, T const& track) + { + float inputFeature = 0.; + switch (idx) { + CHECK_AND_FILL_TRACK(tpcInnerParam); + CHECK_AND_FILL_TRACK(tpcNClsFound); + CHECK_AND_FILL_TRACK(tpcChi2NCl); + CHECK_AND_FILL_TRACK(tpcNSigmaEl); + CHECK_AND_FILL_TRACK(tofNSigmaEl); + CHECK_AND_FILL_TRACK(meanClusterSizeITSobCosTgl); + } + return inputFeature; + } + + /// Method to get the input features vector needed for ML inference + /// \param track is the single track, \param collision is the collision + /// \return inputFeatures vector + template + std::vector getInputFeatures(T const& track) + { + std::vector inputFeatures; + for (const auto& idx : MlResponse::mCachedIndices) { + float inputFeature = return_feature(idx, track); + inputFeatures.emplace_back(inputFeature); + } + return inputFeatures; + } + + /// Method to get the value of variable chosen for binning + /// \param track is the single track, \param collision is the collision + /// \return binning variable + template + float getBinningFeature(T const& track) + { + return return_feature(mCachedIndexBinning, track); + } + + void cacheBinningIndex(std::string const& cfgBinningFeature) + { + setAvailableInputFeatures(); + if (MlResponse::mAvailableInputFeatures.count(cfgBinningFeature)) { + mCachedIndexBinning = MlResponse::mAvailableInputFeatures[cfgBinningFeature]; + } else { + LOG(fatal) << "Binning feature " << cfgBinningFeature << " not available! Please check your configurables."; + } + } + + protected: + /// Method to fill the map of available input features + void setAvailableInputFeatures() + { + MlResponse::mAvailableInputFeatures = { + FILL_MAP_TRACK(tpcInnerParam), + FILL_MAP_TRACK(tpcNClsFound), + FILL_MAP_TRACK(tpcChi2NCl), + FILL_MAP_TRACK(tpcNSigmaEl), + FILL_MAP_TRACK(tofNSigmaEl), + FILL_MAP_TRACK(meanClusterSizeITSobCosTgl), + }; + } + + uint8_t mCachedIndexBinning; // index correspondance between configurable and available input features +}; + +} // namespace o2::analysis + +#undef FILL_MAP_TRACK +#undef CHECK_AND_FILL_TRACK + +#endif // PWGEM_DILEPTON_UTILS_MLRESPONSEPID_H_ diff --git a/PWGEM/Dilepton/Utils/MlResponseSCT.h b/PWGEM/Dilepton/Utils/MlResponseSCT.h new file mode 100644 index 00000000000..71037655e85 --- /dev/null +++ b/PWGEM/Dilepton/Utils/MlResponseSCT.h @@ -0,0 +1,264 @@ +// 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 MlResponseSCT.h +/// \brief Class to compute the ML response for fwdtracks +/// \author Daiki Sekihata + +#ifndef PWGEM_DILEPTON_UTILS_MLRESPONSESCT_H_ +#define PWGEM_DILEPTON_UTILS_MLRESPONSESCT_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_TRACK(FEATURE) \ + { \ + #FEATURE, static_cast(InputFeaturesSCT::FEATURE)} + +// Check if the index of mCachedIndices (index associated to a FEATURE) +// matches the entry in EnumInputFeatures associated to this FEATURE +// if so, the inputFeatures vector is filled with the FEATURE's value +// by calling the corresponding GETTER=FEATURE from track +#define CHECK_AND_FILL_TRACK(GETTER) \ + case static_cast(InputFeaturesSCT::GETTER): { \ + inputFeature = track.GETTER; \ + break; \ + } + +namespace o2::analysis +{ +// possible input features for ML +enum class InputFeaturesSCT : uint8_t { + ptL, + impParXYL, + impParZL, + impPar3DL, + impParXYLinSigma, + impParZLinSigma, + impPar3DLinSigma, + ptH, + tpcNSigmaKa, + impParXYH, + impParZH, + impPar3DH, + impParXYHinSigma, + impParZHinSigma, + impPar3DHinSigma, + signLH, + dcaLH, + massLH, + ptLH, + signedMassLH, + cpa, + cpaXY, + impParXY, + impParZ, + impPar3D, + impParXYinSigma, + impParZinSigma, + impPar3DinSigma, + decayLengthXY, + decayLengthZ, + decayLength3D, + decayLengthXYinSigma, + decayLengthZinSigma, + decayLength3DinSigma, +}; + +namespace pwgem::dilepton::sct +{ +struct candidate { + // lepton info. (don't use them to avoid bias to single lepton) + float ptL{0}; + float impParXYL{0}; + float impParZL{0}; + float impPar3DL{0}; + float impParXYLinSigma{0}; + float impParZLinSigma{0}; + float impPar3DLinSigma{0}; + + // hadron information + float ptH{0}; + float tpcNSigmaKa{0}; + float impParXYH{0}; + float impParZH{0}; + float impPar3DH{0}; + float impParXYHinSigma{0}; + float impParZHinSigma{0}; + float impPar3DHinSigma{0}; + + // LH pair information + int signLH{0}; + float dcaLH{0}; + float massLH{0}; + float ptLH{0}; + float signedMassLH{0}; + float cpa{0}; + float cpaXY{0}; + float impParXY{0}; + float impParZ{0}; + float impPar3D{0}; + float impParXYinSigma{0}; + float impParZinSigma{0}; + float impPar3DinSigma{0}; + float decayLengthXY{0}; + float decayLengthZ{0}; + float decayLength3D{0}; + float decayLengthXYinSigma{0}; + float decayLengthZinSigma{0}; + float decayLength3DinSigma{0}; +}; +} // namespace pwgem::dilepton + +template +class MlResponseSCT : public MlResponse +{ + public: + /// Default constructor + MlResponseSCT() = default; + /// Default destructor + virtual ~MlResponseSCT() = default; + + template + float return_feature(uint8_t idx, T const& track) + { + float inputFeature = 0.; + switch (idx) { + CHECK_AND_FILL_TRACK(ptL); + CHECK_AND_FILL_TRACK(impParXYL); + CHECK_AND_FILL_TRACK(impParZL); + CHECK_AND_FILL_TRACK(impPar3DL); + CHECK_AND_FILL_TRACK(impParXYLinSigma); + CHECK_AND_FILL_TRACK(impParZLinSigma); + CHECK_AND_FILL_TRACK(impPar3DLinSigma); + CHECK_AND_FILL_TRACK(ptH); + CHECK_AND_FILL_TRACK(tpcNSigmaKa); + CHECK_AND_FILL_TRACK(impParXYH); + CHECK_AND_FILL_TRACK(impParZH); + CHECK_AND_FILL_TRACK(impPar3DH); + CHECK_AND_FILL_TRACK(impParXYHinSigma); + CHECK_AND_FILL_TRACK(impParZHinSigma); + CHECK_AND_FILL_TRACK(impPar3DHinSigma); + CHECK_AND_FILL_TRACK(signLH); + CHECK_AND_FILL_TRACK(dcaLH); + CHECK_AND_FILL_TRACK(massLH); + CHECK_AND_FILL_TRACK(ptLH); + CHECK_AND_FILL_TRACK(signedMassLH); + CHECK_AND_FILL_TRACK(cpa); + CHECK_AND_FILL_TRACK(cpaXY); + CHECK_AND_FILL_TRACK(impParXY); + CHECK_AND_FILL_TRACK(impParZ); + CHECK_AND_FILL_TRACK(impPar3D); + CHECK_AND_FILL_TRACK(impParXYinSigma); + CHECK_AND_FILL_TRACK(impParZinSigma); + CHECK_AND_FILL_TRACK(impPar3DinSigma); + CHECK_AND_FILL_TRACK(decayLengthXY); + CHECK_AND_FILL_TRACK(decayLengthZ); + CHECK_AND_FILL_TRACK(decayLength3D); + CHECK_AND_FILL_TRACK(decayLengthXYinSigma); + CHECK_AND_FILL_TRACK(decayLengthZinSigma); + CHECK_AND_FILL_TRACK(decayLength3DinSigma); + } + return inputFeature; + } + + /// Method to get the input features vector needed for ML inference + /// \param track is the single track, \param collision is the collision + /// \return inputFeatures vector + template + std::vector getInputFeatures(T const& track) + { + std::vector inputFeatures; + for (const auto& idx : MlResponse::mCachedIndices) { + float inputFeature = return_feature(idx, track); + inputFeatures.emplace_back(inputFeature); + } + return inputFeatures; + } + + /// Method to get the value of variable chosen for binning + /// \param track is the single track, \param collision is the collision + /// \return binning variable + template + float getBinningFeature(T const& track) + { + return return_feature(mCachedIndexBinning, track); + } + + void cacheBinningIndex(std::string const& cfgBinningFeature) + { + setAvailableInputFeatures(); + if (MlResponse::mAvailableInputFeatures.count(cfgBinningFeature)) { + mCachedIndexBinning = MlResponse::mAvailableInputFeatures[cfgBinningFeature]; + } else { + LOG(fatal) << "Binning feature " << cfgBinningFeature << " not available! Please check your configurables."; + } + } + + protected: + /// Method to fill the map of available input features + void setAvailableInputFeatures() + { + MlResponse::mAvailableInputFeatures = { + FILL_MAP_TRACK(ptL), + FILL_MAP_TRACK(impParXYL), + FILL_MAP_TRACK(impParZL), + FILL_MAP_TRACK(impPar3DL), + FILL_MAP_TRACK(impParXYLinSigma), + FILL_MAP_TRACK(impParZLinSigma), + FILL_MAP_TRACK(impPar3DLinSigma), + FILL_MAP_TRACK(ptH), + FILL_MAP_TRACK(tpcNSigmaKa), + FILL_MAP_TRACK(impParXYH), + FILL_MAP_TRACK(impParZH), + FILL_MAP_TRACK(impPar3DH), + FILL_MAP_TRACK(impParXYHinSigma), + FILL_MAP_TRACK(impParZHinSigma), + FILL_MAP_TRACK(impPar3DHinSigma), + FILL_MAP_TRACK(signLH), + FILL_MAP_TRACK(dcaLH), + FILL_MAP_TRACK(massLH), + FILL_MAP_TRACK(ptLH), + FILL_MAP_TRACK(signedMassLH), + FILL_MAP_TRACK(cpa), + FILL_MAP_TRACK(cpaXY), + FILL_MAP_TRACK(impParXY), + FILL_MAP_TRACK(impParZ), + FILL_MAP_TRACK(impPar3D), + FILL_MAP_TRACK(impParXYinSigma), + FILL_MAP_TRACK(impParZinSigma), + FILL_MAP_TRACK(impPar3DinSigma), + FILL_MAP_TRACK(decayLengthXY), + FILL_MAP_TRACK(decayLengthZ), + FILL_MAP_TRACK(decayLength3D), + FILL_MAP_TRACK(decayLengthXYinSigma), + FILL_MAP_TRACK(decayLengthZinSigma), + FILL_MAP_TRACK(decayLength3DinSigma), + }; + } + + uint8_t mCachedIndexBinning; // index correspondance between configurable and available input features +}; + +} // namespace o2::analysis + +#undef FILL_MAP_TRACK +#undef CHECK_AND_FILL_TRACK + +#endif // PWGEM_DILEPTON_UTILS_MLRESPONSESCT_H_ From 03059ca38949c9f2628d61f3292504312fd6d4ff Mon Sep 17 00:00:00 2001 From: ALICE Action Bot Date: Wed, 20 May 2026 17:04:48 +0000 Subject: [PATCH 2/5] Please consider the following formatting changes --- PWGEM/Dilepton/Utils/MlResponsePID.h | 2 +- PWGEM/Dilepton/Utils/MlResponseSCT.h | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/PWGEM/Dilepton/Utils/MlResponsePID.h b/PWGEM/Dilepton/Utils/MlResponsePID.h index 63e85a6b063..13e8048f798 100644 --- a/PWGEM/Dilepton/Utils/MlResponsePID.h +++ b/PWGEM/Dilepton/Utils/MlResponsePID.h @@ -63,7 +63,7 @@ struct candidate { float tofNSigmaEl{0}; float meanClusterSizeITSobCosTgl{0}; }; -} // namespace pwgem::dilepton +} // namespace pwgem::dilepton::mlpid template class MlResponsePID : public MlResponse diff --git a/PWGEM/Dilepton/Utils/MlResponseSCT.h b/PWGEM/Dilepton/Utils/MlResponseSCT.h index 71037655e85..d0c1353e443 100644 --- a/PWGEM/Dilepton/Utils/MlResponseSCT.h +++ b/PWGEM/Dilepton/Utils/MlResponseSCT.h @@ -124,7 +124,7 @@ struct candidate { float decayLengthZinSigma{0}; float decayLength3DinSigma{0}; }; -} // namespace pwgem::dilepton +} // namespace pwgem::dilepton::sct template class MlResponseSCT : public MlResponse From c03142ec3743fbcafd97dda121a5d18e464cc479 Mon Sep 17 00:00:00 2001 From: Daiki Sekihata Date: Wed, 20 May 2026 19:13:35 +0200 Subject: [PATCH 3/5] Add map header to taggingHFE.cxx --- PWGEM/Dilepton/Tasks/taggingHFE.cxx | 1 + 1 file changed, 1 insertion(+) diff --git a/PWGEM/Dilepton/Tasks/taggingHFE.cxx b/PWGEM/Dilepton/Tasks/taggingHFE.cxx index 163b6b3089e..24d3d72165a 100644 --- a/PWGEM/Dilepton/Tasks/taggingHFE.cxx +++ b/PWGEM/Dilepton/Tasks/taggingHFE.cxx @@ -55,6 +55,7 @@ #include #include #include +#include #include #include #include From 1f50a751ea59f657f73eaaf05c42757bd87b2e03 Mon Sep 17 00:00:00 2001 From: Daiki Sekihata Date: Wed, 20 May 2026 19:22:59 +0200 Subject: [PATCH 4/5] Refactor processWithTTCA and processWithoutTTCA signatures --- PWGEM/Dilepton/Utils/ElectronModule.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/PWGEM/Dilepton/Utils/ElectronModule.h b/PWGEM/Dilepton/Utils/ElectronModule.h index 9abd547c51e..f664fbb4122 100644 --- a/PWGEM/Dilepton/Utils/ElectronModule.h +++ b/PWGEM/Dilepton/Utils/ElectronModule.h @@ -1344,7 +1344,7 @@ class ElectronModule } template - void processWithTTCA(TBCs const& bcs, TCollisions const& collisions, TTracks const& tracks, TV0s const& v0s, TCascades const& cascades, TTrackAssoc const& trackIndices, TMCCollisions const& mcCollisions, TMCParticles const& mcParticles, TProducts& products, THistoregistry& registry, TSliceCache& cache, TPresliceTrack const& perColTrack, TPresliceTrackAssoc const& trackIndicesPerCollision, TPresliceV0 const& perColV0, TPresliceCascade const& perColCasc) + void processWithTTCA(TBCs const& bcs, TCollisions const& collisions, TTracks const& tracks, TV0s const& v0s, TCascades const& cascades, TTrackAssoc const& trackIndices, TMCCollisions const&, TMCParticles const&, TProducts& products, THistoregistry& registry, TSliceCache& cache, TPresliceTrack const& perColTrack, TPresliceTrackAssoc const& trackIndicesPerCollision, TPresliceV0 const& perColV0, TPresliceCascade const& perColCasc) { initCCDB(bcs.begin()); @@ -2106,7 +2106,7 @@ class ElectronModule } template - void processWithoutTTCA(TBCs const& bcs, TCollisions const& collisions, TTracks const& tracks, TV0s const& v0s, TCascades const& cascades, TMCCollisions const& mcCollisions, TMCParticles const& mcParticles, TProducts& products, THistoregistry& registry) + void processWithoutTTCA(TBCs const&, TCollisions const& collisions, TTracks const&, TV0s const&, TCascades const&, TMCCollisions const&, TMCParticles const&, TProducts&, THistoregistry&) { LOGF(info, "processWithoutTTCA is not supported. Bye."); return; From 93cc74d737bb7ffdec742747ab8a1192d10e041b Mon Sep 17 00:00:00 2001 From: Daiki Sekihata Date: Wed, 20 May 2026 19:51:50 +0200 Subject: [PATCH 5/5] Update processWithoutTTCA parameter name for clarity --- PWGEM/Dilepton/Utils/ElectronModule.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/PWGEM/Dilepton/Utils/ElectronModule.h b/PWGEM/Dilepton/Utils/ElectronModule.h index f664fbb4122..87f4de166dc 100644 --- a/PWGEM/Dilepton/Utils/ElectronModule.h +++ b/PWGEM/Dilepton/Utils/ElectronModule.h @@ -2106,7 +2106,7 @@ class ElectronModule } template - void processWithoutTTCA(TBCs const&, TCollisions const& collisions, TTracks const&, TV0s const&, TCascades const&, TMCCollisions const&, TMCParticles const&, TProducts&, THistoregistry&) + void processWithoutTTCA(TBCs const& bcs, TCollisions const& collisions, TTracks const&, TV0s const&, TCascades const&, TMCCollisions const&, TMCParticles const&, TProducts&, THistoregistry&) { LOGF(info, "processWithoutTTCA is not supported. Bye."); return;