Skip to content

Commit 2039b70

Browse files
dsekihatdsekihat
committed
PWGEM/Dilepton: update PID ML in TableProducer
1 parent 9c40da5 commit 2039b70

File tree

2 files changed

+33
-48
lines changed

2 files changed

+33
-48
lines changed

‎PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx‎

Lines changed: 17 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -370,33 +370,15 @@ struct skimmerPrimaryElectron {
370370
return true;
371371
}
372372

373-
template <typename TCollision, typename TTrack>
374-
bool isElectron(TCollision const& collision, TTrack const& track)
373+
template <typename TTrack>
374+
bool isElectron(TTrack const& track)
375375
{
376376
if (includeITSsa && (track.hasITS() && !track.hasTPC() && !track.hasTRD() && !track.hasTOF())) {
377377
return true;
378378
}
379379

380380
if (usePIDML) {
381-
if (track.tpcNSigmaEl() < minTPCNsigmaEl || maxTPCNsigmaEl < track.tpcNSigmaEl()) {
382-
return false;
383-
}
384-
if (track.hasTOF() && (maxTOFNsigmaEl < std::fabs(track.tofNSigmaEl()))) {
385-
return false;
386-
}
387-
388-
// return false;
389-
o2::dataformats::DCA mDcaInfoCov;
390-
mDcaInfoCov.set(999, 999, 999, 999, 999);
391-
auto trackParCov = getTrackParCov(track);
392-
trackParCov.setPID(o2::track::PID::Electron);
393-
mVtx.setPos({collision.posX(), collision.posY(), collision.posZ()});
394-
mVtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ());
395-
o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov);
396-
397-
std::vector<float> inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision);
398-
float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision);
399-
return mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature);
381+
return true;
400382
} else {
401383
return isElectron_TPChadrej(track) || isElectron_TOFreq(track);
402384
}
@@ -457,6 +439,13 @@ struct skimmerPrimaryElectron {
457439
mcTunedTPCSignal = track.mcTunedTPCSignal();
458440
}
459441

442+
float probaEl = 1.0;
443+
if (usePIDML) {
444+
std::vector<float> inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision);
445+
float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision);
446+
probaEl = mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature);
447+
}
448+
460449
emprimaryelectrons(collision.globalIndex(), track.globalIndex(), track.sign(),
461450
pt_recalc, eta_recalc, phi_recalc,
462451
dcaXY, dcaZ, trackParCov.getSigmaY2(), trackParCov.getSigmaZY(), trackParCov.getSigmaZ2(),
@@ -467,7 +456,7 @@ struct skimmerPrimaryElectron {
467456
track.itsClusterSizes(),
468457
track.itsChi2NCl(), track.tofChi2(), track.detectorMap(),
469458
// trackParCov.getTgl(),
470-
isAssociatedToMPC, false, 1.f, mcTunedTPCSignal);
459+
isAssociatedToMPC, false, probaEl, mcTunedTPCSignal);
471460

472461
emprimaryelectronscov(
473462
trackParCov.getX(),
@@ -584,7 +573,7 @@ struct skimmerPrimaryElectron {
584573

585574
auto tracks_per_coll = tracks.sliceBy(perCol, collision.globalIndex());
586575
for (const auto& track : tracks_per_coll) {
587-
if (!checkTrack<false>(collision, track) || !isElectron(collision, track)) {
576+
if (!checkTrack<false>(collision, track) || !isElectron(track)) {
588577
continue;
589578
}
590579
fillTrackTable<false>(collision, track);
@@ -613,7 +602,7 @@ struct skimmerPrimaryElectron {
613602

614603
for (const auto& trackId : trackIdsThisCollision) {
615604
auto track = trackId.template track_as<MyTracks>();
616-
if (!checkTrack<false>(collision, track) || !isElectron(collision, track)) {
605+
if (!checkTrack<false>(collision, track) || !isElectron(track)) {
617606
continue;
618607
}
619608
fillTrackTable<false>(collision, track);
@@ -643,7 +632,7 @@ struct skimmerPrimaryElectron {
643632

644633
auto tracks_per_coll = tracks.sliceBy(perCol, collision.globalIndex());
645634
for (const auto& track : tracks_per_coll) {
646-
if (!checkTrack<false>(collision, track) || !isElectron(collision, track)) {
635+
if (!checkTrack<false>(collision, track) || !isElectron(track)) {
647636
continue;
648637
}
649638
fillTrackTable<false>(collision, track);
@@ -675,7 +664,7 @@ struct skimmerPrimaryElectron {
675664

676665
for (const auto& trackId : trackIdsThisCollision) {
677666
auto track = trackId.template track_as<MyTracks>();
678-
if (!checkTrack<false>(collision, track) || !isElectron(collision, track)) {
667+
if (!checkTrack<false>(collision, track) || !isElectron(track)) {
679668
continue;
680669
}
681670
fillTrackTable<false>(collision, track);
@@ -709,7 +698,7 @@ struct skimmerPrimaryElectron {
709698

710699
auto tracks_per_coll = tracks.sliceBy(perCol, collision.globalIndex());
711700
for (const auto& track : tracks_per_coll) {
712-
if (!checkTrack<true>(collision, track) || !isElectron(collision, track)) {
701+
if (!checkTrack<true>(collision, track) || !isElectron(track)) {
713702
continue;
714703
}
715704
fillTrackTable<true>(collision, track);
@@ -740,7 +729,7 @@ struct skimmerPrimaryElectron {
740729

741730
for (const auto& trackId : trackIdsThisCollision) {
742731
auto track = trackId.template track_as<MyTracksMC>();
743-
if (!checkTrack<true>(collision, track) || !isElectron(collision, track)) {
732+
if (!checkTrack<true>(collision, track) || !isElectron(track)) {
744733
continue;
745734
}
746735
fillTrackTable<true>(collision, track);

‎PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronQC.cxx‎

Lines changed: 16 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -420,26 +420,15 @@ struct skimmerPrimaryElectronQC {
420420
template <typename TTrack>
421421
bool isElectron(TTrack const& track)
422422
{
423-
if (track.hasTPC() && (trackcut.minTPCNsigmaEl < track.tpcNSigmaEl() && track.tpcNSigmaEl() < trackcut.maxTPCNsigmaEl)) {
424-
return true; // accept ITSsa too
425-
} else {
426-
return false;
423+
if (track.hasTPC()) {
424+
if (trackcut.minTPCNsigmaEl < track.tpcNSigmaEl() && track.tpcNSigmaEl() < trackcut.maxTPCNsigmaEl) {
425+
return true;
426+
} else {
427+
return false;
428+
}
429+
} else { // accept ITSsa too
430+
return true;
427431
}
428-
429-
// if (usePIDML) {
430-
// mDcaInfoCov.set(999, 999, 999, 999, 999);
431-
// auto trackParCov = getTrackParCov(track);
432-
// trackParCov.setPID(o2::track::PID::Electron);
433-
// mVtx.setPos({collision.posX(), collision.posY(), collision.posZ()});
434-
// mVtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ());
435-
// o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov);
436-
437-
// std::vector<float> inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision);
438-
// float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision);
439-
// return mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature);
440-
// } else {
441-
// return isElectronTPC(track);
442-
// }
443432
}
444433

445434
template <typename TTrack>
@@ -474,6 +463,13 @@ struct skimmerPrimaryElectronQC {
474463
mcTunedTPCSignal = track.mcTunedTPCSignal();
475464
}
476465

466+
float probaEl = 1.0;
467+
if (usePIDML) {
468+
std::vector<float> inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision);
469+
float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision);
470+
probaEl = mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature);
471+
}
472+
477473
emprimaryelectrons(collision.globalIndex(), track.globalIndex(), track.sign(),
478474
pt_recalc, eta_recalc, phi_recalc,
479475
dcaXY, dcaZ, trackParCov.getSigmaY2(), trackParCov.getSigmaZY(), trackParCov.getSigmaZ2(),
@@ -484,7 +480,7 @@ struct skimmerPrimaryElectronQC {
484480
track.itsClusterSizes(),
485481
track.itsChi2NCl(), track.tofChi2(), track.detectorMap(),
486482
// trackParCov.getTgl(),
487-
isAssociatedToMPC, false, 1.f, mcTunedTPCSignal);
483+
isAssociatedToMPC, false, probaEl, mcTunedTPCSignal);
488484

489485
emprimaryelectronscov(
490486
trackParCov.getX(),

0 commit comments

Comments
 (0)