From 2039b70b476dc1fa2f21d2f67c0edd7d3e489028 Mon Sep 17 00:00:00 2001 From: Daiki Sekihata Date: Thu, 21 Aug 2025 00:20:01 +0200 Subject: [PATCH] PWGEM/Dilepton: update PID ML in TableProducer --- .../TableProducer/skimmerPrimaryElectron.cxx | 45 +++++++------------ .../skimmerPrimaryElectronQC.cxx | 36 +++++++-------- 2 files changed, 33 insertions(+), 48 deletions(-) diff --git a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx index 778c765aebd..bf1e2b61505 100644 --- a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx +++ b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx @@ -370,33 +370,15 @@ struct skimmerPrimaryElectron { return true; } - template - bool isElectron(TCollision const& collision, TTrack const& track) + template + bool isElectron(TTrack const& track) { if (includeITSsa && (track.hasITS() && !track.hasTPC() && !track.hasTRD() && !track.hasTOF())) { return true; } if (usePIDML) { - if (track.tpcNSigmaEl() < minTPCNsigmaEl || maxTPCNsigmaEl < track.tpcNSigmaEl()) { - return false; - } - if (track.hasTOF() && (maxTOFNsigmaEl < std::fabs(track.tofNSigmaEl()))) { - return false; - } - - // return false; - o2::dataformats::DCA mDcaInfoCov; - mDcaInfoCov.set(999, 999, 999, 999, 999); - auto trackParCov = getTrackParCov(track); - trackParCov.setPID(o2::track::PID::Electron); - mVtx.setPos({collision.posX(), collision.posY(), collision.posZ()}); - mVtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ()); - o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov); - - std::vector inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision); - float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision); - return mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature); + return true; } else { return isElectron_TPChadrej(track) || isElectron_TOFreq(track); } @@ -457,6 +439,13 @@ struct skimmerPrimaryElectron { mcTunedTPCSignal = track.mcTunedTPCSignal(); } + float probaEl = 1.0; + if (usePIDML) { + std::vector inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision); + float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision); + probaEl = mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature); + } + emprimaryelectrons(collision.globalIndex(), track.globalIndex(), track.sign(), pt_recalc, eta_recalc, phi_recalc, dcaXY, dcaZ, trackParCov.getSigmaY2(), trackParCov.getSigmaZY(), trackParCov.getSigmaZ2(), @@ -467,7 +456,7 @@ struct skimmerPrimaryElectron { track.itsClusterSizes(), track.itsChi2NCl(), track.tofChi2(), track.detectorMap(), // trackParCov.getTgl(), - isAssociatedToMPC, false, 1.f, mcTunedTPCSignal); + isAssociatedToMPC, false, probaEl, mcTunedTPCSignal); emprimaryelectronscov( trackParCov.getX(), @@ -584,7 +573,7 @@ struct skimmerPrimaryElectron { auto tracks_per_coll = tracks.sliceBy(perCol, collision.globalIndex()); for (const auto& track : tracks_per_coll) { - if (!checkTrack(collision, track) || !isElectron(collision, track)) { + if (!checkTrack(collision, track) || !isElectron(track)) { continue; } fillTrackTable(collision, track); @@ -613,7 +602,7 @@ struct skimmerPrimaryElectron { for (const auto& trackId : trackIdsThisCollision) { auto track = trackId.template track_as(); - if (!checkTrack(collision, track) || !isElectron(collision, track)) { + if (!checkTrack(collision, track) || !isElectron(track)) { continue; } fillTrackTable(collision, track); @@ -643,7 +632,7 @@ struct skimmerPrimaryElectron { auto tracks_per_coll = tracks.sliceBy(perCol, collision.globalIndex()); for (const auto& track : tracks_per_coll) { - if (!checkTrack(collision, track) || !isElectron(collision, track)) { + if (!checkTrack(collision, track) || !isElectron(track)) { continue; } fillTrackTable(collision, track); @@ -675,7 +664,7 @@ struct skimmerPrimaryElectron { for (const auto& trackId : trackIdsThisCollision) { auto track = trackId.template track_as(); - if (!checkTrack(collision, track) || !isElectron(collision, track)) { + if (!checkTrack(collision, track) || !isElectron(track)) { continue; } fillTrackTable(collision, track); @@ -709,7 +698,7 @@ struct skimmerPrimaryElectron { auto tracks_per_coll = tracks.sliceBy(perCol, collision.globalIndex()); for (const auto& track : tracks_per_coll) { - if (!checkTrack(collision, track) || !isElectron(collision, track)) { + if (!checkTrack(collision, track) || !isElectron(track)) { continue; } fillTrackTable(collision, track); @@ -740,7 +729,7 @@ struct skimmerPrimaryElectron { for (const auto& trackId : trackIdsThisCollision) { auto track = trackId.template track_as(); - if (!checkTrack(collision, track) || !isElectron(collision, track)) { + if (!checkTrack(collision, track) || !isElectron(track)) { continue; } fillTrackTable(collision, track); diff --git a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronQC.cxx b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronQC.cxx index c448464ceab..68b0fad8593 100644 --- a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronQC.cxx +++ b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronQC.cxx @@ -420,26 +420,15 @@ struct skimmerPrimaryElectronQC { template bool isElectron(TTrack const& track) { - if (track.hasTPC() && (trackcut.minTPCNsigmaEl < track.tpcNSigmaEl() && track.tpcNSigmaEl() < trackcut.maxTPCNsigmaEl)) { - return true; // accept ITSsa too - } else { - return false; + if (track.hasTPC()) { + if (trackcut.minTPCNsigmaEl < track.tpcNSigmaEl() && track.tpcNSigmaEl() < trackcut.maxTPCNsigmaEl) { + return true; + } else { + return false; + } + } else { // accept ITSsa too + return true; } - - // if (usePIDML) { - // mDcaInfoCov.set(999, 999, 999, 999, 999); - // auto trackParCov = getTrackParCov(track); - // trackParCov.setPID(o2::track::PID::Electron); - // mVtx.setPos({collision.posX(), collision.posY(), collision.posZ()}); - // mVtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ()); - // o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov); - - // std::vector inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision); - // float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision); - // return mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature); - // } else { - // return isElectronTPC(track); - // } } template @@ -474,6 +463,13 @@ struct skimmerPrimaryElectronQC { mcTunedTPCSignal = track.mcTunedTPCSignal(); } + float probaEl = 1.0; + if (usePIDML) { + std::vector inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision); + float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision); + probaEl = mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature); + } + emprimaryelectrons(collision.globalIndex(), track.globalIndex(), track.sign(), pt_recalc, eta_recalc, phi_recalc, dcaXY, dcaZ, trackParCov.getSigmaY2(), trackParCov.getSigmaZY(), trackParCov.getSigmaZ2(), @@ -484,7 +480,7 @@ struct skimmerPrimaryElectronQC { track.itsClusterSizes(), track.itsChi2NCl(), track.tofChi2(), track.detectorMap(), // trackParCov.getTgl(), - isAssociatedToMPC, false, 1.f, mcTunedTPCSignal); + isAssociatedToMPC, false, probaEl, mcTunedTPCSignal); emprimaryelectronscov( trackParCov.getX(),