diff --git a/PWGEM/Dilepton/DataModel/lmeeMLTables.h b/PWGEM/Dilepton/DataModel/lmeeMLTables.h index d2f750ac2fc..eb4fdee60b2 100644 --- a/PWGEM/Dilepton/DataModel/lmeeMLTables.h +++ b/PWGEM/Dilepton/DataModel/lmeeMLTables.h @@ -162,6 +162,8 @@ DECLARE_SOA_COLUMN(IsPrimaryMCHMID, isPrimaryMCHMID, bool); //! DECLARE_SOA_COLUMN(IsCorrectMatch, isCorrectMatch, bool); //! DECLARE_SOA_COLUMN(PdgCodeMFT, pdgCodeMFT, int); //! DECLARE_SOA_COLUMN(PdgCodeMCHMID, pdgCodeMCHMID, int); //! +DECLARE_SOA_COLUMN(MatchMCHTrackId, mchTrackId, int); //! +DECLARE_SOA_COLUMN(DFId, dfId, uint64_t); //! DECLARE_SOA_COLUMN(MultMFT, multMFT, uint16_t); //! number of MFTsa tracks per collision } // namespace emmlfwdtrack @@ -178,7 +180,7 @@ DECLARE_SOA_TABLE(EMFwdTracksForML, "AOD", "EMFWDTRKML", //! fwdtrack::Chi2MatchMCHMFT, emmlfwdtrack::PdgCodeMFT, emmlfwdtrack::IsPrimaryMFT, emmlfwdtrack::PdgCodeMCHMID, emmlfwdtrack::IsPrimaryMCHMID, - emmlfwdtrack::IsCorrectMatch); + emmlfwdtrack::IsCorrectMatch, emmlfwdtrack::MatchMCHTrackId, emmlfwdtrack::DFId); // iterators using EMFwdTrackForML = EMFwdTracksForML::iterator; diff --git a/PWGEM/Dilepton/TableProducer/treeCreatorMuonML.cxx b/PWGEM/Dilepton/TableProducer/treeCreatorMuonML.cxx index ece00db378a..05eaeaf9ee1 100644 --- a/PWGEM/Dilepton/TableProducer/treeCreatorMuonML.cxx +++ b/PWGEM/Dilepton/TableProducer/treeCreatorMuonML.cxx @@ -266,8 +266,10 @@ struct TreeCreatorMuonML { std::unordered_map map_mfttrackcovs; - void processWithMFTCov(MyCollisionsMC const& collisions, aod::BCsWithTimestamps const&, MyFwdTracksMC const&, MyMFTTracksMC const& mfttracks, aod::MFTTracksCov const& mftCovs, aod::McParticles const&, aod::McCollisions const&) + void processWithMFTCov(MyCollisionsMC const& collisions, aod::BCsWithTimestamps const&, MyFwdTracksMC const&, MyMFTTracksMC const& mfttracks, aod::MFTTracksCov const& mftCovs, aod::McParticles const&, aod::McCollisions const&, aod::Origins const& origins) { + uint64_t dfId = origins.iteratorAt(0).dataframeID(); + for (const auto& mfttrackCov : mftCovs) { map_mfttrackcovs[mfttrackCov.matchMFTTrackId()] = mfttrackCov.globalIndex(); } @@ -378,7 +380,7 @@ struct TreeCreatorMuonML { fwdtrack.chi2MatchMCHMFT(), mcParticle_MFT.pdgCode(), isPrimary_MFT, mcParticle_MCHMID.pdgCode(), isPrimary_MCHMID, - isMatched); + isMatched, mchtrack.globalIndex(), dfId); trackErrTable(tglErrMFTatMP, phiErrMFTatMP, xErrMFTatMP, yErrMFTatMP, tglErrMCHMIDatMP, phiErrMCHMIDatMP, xErrMCHMIDatMP, yErrMCHMIDatMP); diff --git a/PWGEM/Dilepton/Tasks/taggingHFE.cxx b/PWGEM/Dilepton/Tasks/taggingHFE.cxx index a07ed257816..e04f8bc2244 100644 --- a/PWGEM/Dilepton/Tasks/taggingHFE.cxx +++ b/PWGEM/Dilepton/Tasks/taggingHFE.cxx @@ -55,6 +55,7 @@ #include #include #include +#include #include #include #include @@ -94,6 +95,7 @@ struct taggingHFE { Configurable skipGRPOquery{"skipGRPOquery", true, "skip grpo query"}; Configurable d_bz_input{"d_bz_input", -999, "bz field in kG, -999 is automatic"}; Configurable cfgPdgLepton{"cfgPdgLepton", 11, "pdg code of desired lepton: 11 or 13"}; + Configurable cfgDownSampling{"cfgDownSampling", 1.1, "down sampling for fake matches"}; struct : ConfigurableGroup { std::string prefix = "dcaFitterGroup_eK"; @@ -267,6 +269,10 @@ struct taggingHFE { ccdb->setLocalObjectValidityChecking(); ccdb->setFatalWhenNull(false); + std::random_device seed_gen; + engine = std::mt19937(seed_gen()); + dist01 = std::uniform_real_distribution(0.0f, 1.0f); + fitter_eK.setPropagateToPCA(true); fitter_eK.setMaxR(20.f); fitter_eK.setMinParamChange(1e-3); @@ -303,6 +309,9 @@ struct taggingHFE { int mRunNumber; float d_bz; Service ccdb; + std::mt19937 engine; + std::uniform_real_distribution dist01; + // o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrNONE; o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrLUT; const o2::dataformats::MeanVertexObject* mMeanVtx = nullptr; @@ -970,6 +979,10 @@ struct taggingHFE { mVtx.setPos({collision.posX(), collision.posY(), collision.posZ()}); mVtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ()); + if (dist01(engine) > cfgDownSampling) { // random sampling, if necessary + continue; + } + const auto& trackIdsThisCollision = trackIndices.sliceBy(trackIndicesPerCollision, collision.globalIndex()); electronIds.reserve(trackIdsThisCollision.size()); positronIds.reserve(trackIdsThisCollision.size());