Skip to content

Commit f8d472e

Browse files
kgwizdzikgwizdzi
authored
[PWGCF] FemtoUniverse - producer and D0 task update (ML) (#9562)
1 parent 35d89ce commit f8d472e

File tree

2 files changed

+269
-26
lines changed

2 files changed

+269
-26
lines changed

‎PWGCF/FemtoUniverse/TableProducer/femtoUniverseProducerTask.cxx‎

Lines changed: 163 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,9 @@ struct FemtoUniverseProducerTask {
299299
struct : o2::framework::ConfigurableGroup {
300300
Configurable<float> confD0D0barCandMaxY{"confD0D0barCandMaxY", -1., "max. cand. rapidity"};
301301
Configurable<float> confD0D0barCandEtaCut{"confD0D0barCandEtaCut", 0.8, "max. cand. pseudorapidity"};
302+
Configurable<bool> storeD0D0barDoubleMassHypo{"storeD0D0barDoubleMassHypo", false, "Store D0/D0bar cand. which pass selection criteria for both, D0 and D0bar"};
303+
Configurable<bool> applyMLSelD0D0bar{"applyMLSelD0D0bar", false, "Use ML D0/D0bar selection"};
304+
Configurable<std::vector<int>> classMlD0D0bar{"classMlD0D0bar", {0, 1, 2}, "Indexes of ML scores to be stored. Three indexes max."};
302305
} ConfD0Selection;
303306

304307
HfHelper hfHelper;
@@ -458,10 +461,10 @@ struct FemtoUniverseProducerTask {
458461

459462
void init(InitContext&)
460463
{
461-
if ((doprocessFullData || doprocessTrackPhiData || doprocessTrackData || doprocessTrackV0 || doprocessTrackCascadeData || doprocessTrackD0mesonData || doprocessTrackCentRun2Data || doprocessTrackV0CentRun2Data || doprocessTrackCentRun3Data || doprocessV0CentRun3Data || doprocessCascadeCentRun3Data || doprocessTrackDataCentPP) == false && (doprocessFullMC || doprocessTrackMC || doprocessTrackMCTruth || doprocessTrackMCGen || doprocessTruthAndFullMC || doprocessFullMCCent || doprocessTrackCentRun3DataMC) == false) {
464+
if ((doprocessFullData || doprocessTrackPhiData || doprocessTrackData || doprocessTrackV0 || doprocessTrackCascadeData || doprocessTrackD0mesonData || doprocessTrackD0DataML || doprocessTrackCentRun2Data || doprocessTrackV0CentRun2Data || doprocessTrackCentRun3Data || doprocessV0CentRun3Data || doprocessCascadeCentRun3Data || doprocessTrackDataCentPP) == false && (doprocessFullMC || doprocessTrackMC || doprocessTrackMCTruth || doprocessTrackMCGen || doprocessTruthAndFullMC || doprocessFullMCCent || doprocessTrackCentRun3DataMC) == false) {
462465
LOGF(fatal, "Neither processFullData nor processFullMC enabled. Please choose one.");
463466
}
464-
if ((doprocessFullData || doprocessTrackPhiData || doprocessTrackData || doprocessTrackV0 || doprocessTrackCascadeData || doprocessTrackD0mesonData || doprocessTrackCentRun2Data || doprocessTrackV0CentRun2Data || doprocessTrackCentRun3Data || doprocessV0CentRun3Data || doprocessCascadeCentRun3Data || doprocessTrackDataCentPP) == true && (doprocessFullMC || doprocessTrackMC || doprocessTrackMCTruth || doprocessTrackMCGen || doprocessTruthAndFullMC || doprocessFullMCCent || doprocessTrackCentRun3DataMC) == true) {
467+
if ((doprocessFullData || doprocessTrackPhiData || doprocessTrackData || doprocessTrackV0 || doprocessTrackCascadeData || doprocessTrackD0mesonData || doprocessTrackD0DataML || doprocessTrackCentRun2Data || doprocessTrackV0CentRun2Data || doprocessTrackCentRun3Data || doprocessV0CentRun3Data || doprocessCascadeCentRun3Data || doprocessTrackDataCentPP) == true && (doprocessFullMC || doprocessTrackMC || doprocessTrackMCTruth || doprocessTrackMCGen || doprocessTruthAndFullMC || doprocessFullMCCent || doprocessTrackCentRun3DataMC) == true) {
465468
LOGF(fatal,
466469
"Cannot enable process Data and process MC at the same time. "
467470
"Please choose one.");
@@ -1253,8 +1256,13 @@ struct FemtoUniverseProducerTask {
12531256
} else if (hfCand.isSelD0() == 1 && hfCand.isSelD0bar() == 1) {
12541257
invMassD0 = hfHelper.invMassD0ToPiK(hfCand);
12551258
invMassD0bar = hfHelper.invMassD0barToKPi(hfCand);
1256-
isD0D0bar = true;
1257-
daughFlag = 0;
1259+
if (ConfD0Selection.storeD0D0barDoubleMassHypo) {
1260+
isD0D0bar = true;
1261+
daughFlag = 0;
1262+
} else {
1263+
isD0D0bar = false;
1264+
daughFlag = 0;
1265+
}
12581266
} else {
12591267
invMassD0 = 0.0;
12601268
invMassD0bar = 0.0;
@@ -1308,7 +1316,141 @@ struct FemtoUniverseProducerTask {
13081316
aod::femtouniverseparticle::ParticleType::kD0,
13091317
-999, // cut, CutContainerType
13101318
-999, // PID, CutContainerType
1319+
-999.,
1320+
indexChildID,
1321+
invMassD0, // D0 mass (mLambda)
1322+
invMassD0bar); // D0bar mass (mAntiLambda)
1323+
1324+
if (confIsDebug) {
1325+
fillDebugParticle<false, true, false>(postrack); // QA for positive daughter
1326+
fillDebugParticle<false, true, false>(negtrack); // QA for negative daughter
1327+
fillDebugParticle<false, true, false>(hfCand); // QA for D0/D0bar
1328+
}
1329+
if constexpr (isMC) {
1330+
fillMCParticle(hfCand, o2::aod::femtouniverseparticle::ParticleType::kD0);
1331+
}
1332+
}
1333+
}
1334+
}
1335+
1336+
template <bool isMC, typename HfCandidate, typename TrackType, typename CollisionType>
1337+
void fillD0D0barUsingML(CollisionType const&, TrackType const&, HfCandidate const& hfCands)
1338+
{
1339+
std::vector<int> childIDs = {0, 0}; // these IDs are necessary to keep track of the children
1340+
std::vector<int> tmpIDtrack; // this vector keeps track of the matching of the primary track table row <-> aod::track table global index
1341+
double invMassD0 = 0.0;
1342+
double invMassD0bar = 0.0;
1343+
bool isD0D0bar = false;
1344+
std::vector<float> outputMlD0D0bar = {-1., -1., -1.}; // this vector keeps the probabilities from the ML model for D0/D0bar
1345+
uint8_t daughFlag = 0; // flag = 0 (daugh of D0 or D0bar), 1 (daug of D0), -1 (daugh of D0bar)
1346+
1347+
for (const auto& hfCand : hfCands) {
1348+
1349+
if (!(hfCand.hfflag() & 1 << aod::hf_cand_2prong::DecayType::D0ToPiK)) {
1350+
continue;
1351+
}
1352+
1353+
if (ConfD0Selection.confD0D0barCandMaxY >= 0. && std::abs(hfHelper.yD0(hfCand)) > ConfD0Selection.confD0D0barCandMaxY) {
1354+
continue;
1355+
}
1356+
1357+
if (std::abs(hfCand.eta()) > ConfD0Selection.confD0D0barCandEtaCut) {
1358+
continue;
1359+
}
1360+
1361+
// int postrackID = hfCand.prong0().globalIndex();
1362+
int postrackID = hfCand.prong0Id(); // Index to first prong
1363+
int rowInPrimaryTrackTablePos = -1;
1364+
rowInPrimaryTrackTablePos = getRowDaughters(postrackID, tmpIDtrack);
1365+
childIDs[0] = rowInPrimaryTrackTablePos;
1366+
childIDs[1] = 0;
1367+
auto postrack = hfCand.template prong0_as<TrackType>();
1368+
auto negtrack = hfCand.template prong1_as<TrackType>();
1369+
1370+
if (hfCand.isSelD0() == 1 && hfCand.isSelD0bar() == 0) {
1371+
invMassD0 = hfHelper.invMassD0ToPiK(hfCand);
1372+
invMassD0bar = -hfHelper.invMassD0barToKPi(hfCand);
1373+
if (ConfD0Selection.applyMLSelD0D0bar) {
1374+
for (unsigned int iClass = 0; iClass < ConfD0Selection.classMlD0D0bar->size(); iClass++) {
1375+
outputMlD0D0bar[iClass] = hfCand.mlProbD0()[ConfD0Selection.classMlD0D0bar->at(iClass)];
1376+
}
1377+
}
1378+
isD0D0bar = true;
1379+
daughFlag = 1;
1380+
} else if (hfCand.isSelD0() == 0 && hfCand.isSelD0bar() == 1) {
1381+
invMassD0 = -hfHelper.invMassD0ToPiK(hfCand);
1382+
invMassD0bar = hfHelper.invMassD0barToKPi(hfCand);
1383+
if (ConfD0Selection.applyMLSelD0D0bar) {
1384+
for (unsigned int iClass = 0; iClass < ConfD0Selection.classMlD0D0bar->size(); iClass++) {
1385+
outputMlD0D0bar[iClass] = hfCand.mlProbD0bar()[ConfD0Selection.classMlD0D0bar->at(iClass)];
1386+
}
1387+
}
1388+
isD0D0bar = true;
1389+
daughFlag = -1;
1390+
} else if (hfCand.isSelD0() == 1 && hfCand.isSelD0bar() == 1) {
1391+
invMassD0 = hfHelper.invMassD0ToPiK(hfCand);
1392+
invMassD0bar = hfHelper.invMassD0barToKPi(hfCand);
1393+
if (ConfD0Selection.storeD0D0barDoubleMassHypo) {
1394+
isD0D0bar = true;
1395+
daughFlag = 0;
1396+
} else {
1397+
isD0D0bar = false;
1398+
daughFlag = 0;
1399+
}
1400+
} else {
1401+
invMassD0 = 0.0;
1402+
invMassD0bar = 0.0;
1403+
isD0D0bar = false;
1404+
}
1405+
1406+
if (isD0D0bar) {
1407+
outputParts(outputCollision.lastIndex(),
1408+
hfCand.ptProng0(),
1409+
RecoDecay::eta(std::array{hfCand.pxProng0(), hfCand.pyProng0(), hfCand.pzProng0()}), // eta
1410+
RecoDecay::phi(hfCand.pxProng0(), hfCand.pyProng0()), // phi
1411+
aod::femtouniverseparticle::ParticleType::kD0Child,
1412+
-999, // cutContainerV0.at(femto_universe_v0_selection::V0ContainerPosition::kPosCuts),
1413+
-999, // cutContainerV0.at(femto_universe_v0_selection::V0ContainerPosition::kPosPID),
1414+
-999,
1415+
childIDs,
1416+
postrack.sign(), // D0 mass -> positive daughter of D0/D0bar
1417+
daughFlag); // D0bar mass -> sign that the daugh is from D0 or D0 decay
1418+
const int rowOfPosTrack = outputParts.lastIndex();
1419+
if constexpr (isMC) {
1420+
fillMCParticle(postrack, o2::aod::femtouniverseparticle::ParticleType::kD0Child);
1421+
}
1422+
// int negtrackID = hfCand.prong1().globalIndex();
1423+
int negtrackID = hfCand.prong1Id();
1424+
int rowInPrimaryTrackTableNeg = -1;
1425+
rowInPrimaryTrackTableNeg = getRowDaughters(negtrackID, tmpIDtrack);
1426+
childIDs[0] = 0;
1427+
childIDs[1] = rowInPrimaryTrackTableNeg;
1428+
1429+
outputParts(outputCollision.lastIndex(),
1430+
hfCand.ptProng1(),
1431+
RecoDecay::eta(std::array{hfCand.pxProng1(), hfCand.pyProng1(), hfCand.pzProng1()}), // eta
1432+
RecoDecay::phi(hfCand.pxProng1(), hfCand.pyProng1()), // phi
1433+
aod::femtouniverseparticle::ParticleType::kD0Child,
1434+
-999, // cutContainerV0.at(femto_universe_v0_selection::V0ContainerPosition::kNegCuts),
1435+
-999, // cutContainerV0.at(femto_universe_v0_selection::V0ContainerPosition::kNegPID),
13111436
-999,
1437+
childIDs,
1438+
negtrack.sign(), // negative daughter of D0/D0bar
1439+
daughFlag); // sign that the daugh is from D0 or D0 decay
1440+
const int rowOfNegTrack = outputParts.lastIndex();
1441+
if constexpr (isMC) {
1442+
fillMCParticle(negtrack, o2::aod::femtouniverseparticle::ParticleType::kD0Child);
1443+
}
1444+
std::vector<int> indexChildID = {rowOfPosTrack, rowOfNegTrack};
1445+
1446+
outputParts(outputCollision.lastIndex(),
1447+
hfCand.pt(),
1448+
hfCand.eta(),
1449+
hfCand.phi(),
1450+
aod::femtouniverseparticle::ParticleType::kD0,
1451+
-999, // cut, CutContainerType
1452+
-999, // PID, CutContainerType
1453+
outputMlD0D0bar[0], // saving only the probaility for store class 1 - background
13121454
indexChildID,
13131455
invMassD0, // D0 mass (mLambda)
13141456
invMassD0bar); // D0bar mass (mAntiLambda)
@@ -1783,6 +1925,23 @@ struct FemtoUniverseProducerTask {
17831925
PROCESS_SWITCH(FemtoUniverseProducerTask, processTrackD0mesonData,
17841926
"Provide experimental data for track D0 meson", false);
17851927

1928+
void processTrackD0DataML(aod::FemtoFullCollision const& col,
1929+
aod::BCsWithTimestamps const&,
1930+
soa::Filtered<aod::FemtoFullTracks> const& tracks,
1931+
soa::Join<aod::HfCand2Prong, aod::HfSelD0, aod::HfMlD0> const& candidates)
1932+
{
1933+
// get magnetic field for run
1934+
getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
1935+
// fill the tables
1936+
const auto colcheck = fillCollisions<false>(col, tracks);
1937+
if (colcheck) {
1938+
fillTracks<false>(tracks);
1939+
fillD0D0barUsingML<false>(col, tracks, candidates);
1940+
}
1941+
}
1942+
PROCESS_SWITCH(FemtoUniverseProducerTask, processTrackD0DataML,
1943+
"Provide experimental data for track D0 meson using ML selection for D0s", false);
1944+
17861945
void processTrackMCTruth(aod::McCollision const&,
17871946
soa::SmallGroups<soa::Join<aod::Collisions, aod::EvSels, aod::Mults, aod::McCollisionLabels>> const& collisions,
17881947
aod::McParticles const& mcParticles,

0 commit comments

Comments
 (0)