Skip to content

Commit 0d9a1d6

Browse files
vkuceraalibuild
vkucera
and
alibuild
authored
[PWGHF] Refactor derived-data creators (#8581)
Co-authored-by: ALICE Builder <alibuild@users.noreply.github.com>
1 parent 0e044e9 commit 0d9a1d6

File tree

4 files changed

+335
-532
lines changed

4 files changed

+335
-532
lines changed

‎PWGHF/TableProducer/derivedDataCreatorBplusToD0Pi.cxx‎

Lines changed: 34 additions & 173 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,17 @@ using namespace o2::analysis::hf_derived;
4444

4545
/// Writes the full information in an output TTree
4646
struct HfDerivedDataCreatorBplusToD0Pi {
47+
HfProducesDerivedData<
48+
o2::aod::HfBplusBases,
49+
o2::aod::HfBplusCollBases,
50+
o2::aod::HfBplusCollIds,
51+
o2::aod::HfBplusMcCollBases,
52+
o2::aod::HfBplusMcCollIds,
53+
o2::aod::HfBplusMcRCollIds,
54+
o2::aod::HfBplusPBases,
55+
o2::aod::HfBplusPIds>
56+
rowsCommon;
4757
// Candidates
48-
Produces<o2::aod::HfBplusBases> rowCandidateBase;
4958
Produces<o2::aod::HfBplusPars> rowCandidatePar;
5059
Produces<o2::aod::HfBplusParD0s> rowCandidateParD0;
5160
Produces<o2::aod::HfBplusParEs> rowCandidateParE;
@@ -54,19 +63,9 @@ struct HfDerivedDataCreatorBplusToD0Pi {
5463
Produces<o2::aod::HfBplusMlD0s> rowCandidateMlD0;
5564
Produces<o2::aod::HfBplusIds> rowCandidateId;
5665
Produces<o2::aod::HfBplusMcs> rowCandidateMc;
57-
// Collisions
58-
Produces<o2::aod::HfBplusCollBases> rowCollBase;
59-
Produces<o2::aod::HfBplusCollIds> rowCollId;
60-
// MC collisions
61-
Produces<o2::aod::HfBplusMcCollBases> rowMcCollBase;
62-
Produces<o2::aod::HfBplusMcCollIds> rowMcCollId;
63-
Produces<o2::aod::HfBplusMcRCollIds> rowMcRCollId;
64-
// MC particles
65-
Produces<o2::aod::HfBplusPBases> rowParticleBase;
66-
Produces<o2::aod::HfBplusPIds> rowParticleId;
6766

6867
// Switches for filling tables
69-
Configurable<bool> fillCandidateBase{"fillCandidateBase", true, "Fill candidate base properties"};
68+
HfConfigurableDerivedData confDerData;
7069
Configurable<bool> fillCandidatePar{"fillCandidatePar", true, "Fill candidate parameters"};
7170
Configurable<bool> fillCandidateParD0{"fillCandidateParD0", true, "Fill D0 candidate parameters"};
7271
Configurable<bool> fillCandidateParE{"fillCandidateParE", true, "Fill candidate extended parameters"};
@@ -75,21 +74,13 @@ struct HfDerivedDataCreatorBplusToD0Pi {
7574
Configurable<bool> fillCandidateMlD0{"fillCandidateMlD0", true, "Fill D0 candidate selection ML scores"};
7675
Configurable<bool> fillCandidateId{"fillCandidateId", true, "Fill original indices from the candidate table"};
7776
Configurable<bool> fillCandidateMc{"fillCandidateMc", true, "Fill candidate MC info"};
78-
Configurable<bool> fillCollBase{"fillCollBase", true, "Fill collision base properties"};
79-
Configurable<bool> fillCollId{"fillCollId", true, "Fill original collision indices"};
80-
Configurable<bool> fillMcCollBase{"fillMcCollBase", true, "Fill MC collision base properties"};
81-
Configurable<bool> fillMcCollId{"fillMcCollId", true, "Fill original MC collision indices"};
82-
Configurable<bool> fillMcRCollId{"fillMcRCollId", true, "Fill indices of saved derived reconstructed collisions matched to saved derived MC collisions"};
83-
Configurable<bool> fillParticleBase{"fillParticleBase", true, "Fill MC particle properties"};
84-
Configurable<bool> fillParticleId{"fillParticleId", true, "Fill original MC indices"};
8577
// Parameters for production of training samples
8678
Configurable<float> downSampleBkgFactor{"downSampleBkgFactor", 1., "Fraction of background candidates to keep for ML trainings"};
8779
Configurable<float> ptMaxForDownSample{"ptMaxForDownSample", 10., "Maximum pt for the application of the downsampling factor"};
8880

8981
HfHelper hfHelper;
9082
SliceCache cache;
91-
std::map<int, std::vector<int>> matchedCollisions; // indices of derived reconstructed collisions matched to the global indices of MC collisions
92-
std::map<int, bool> hasMcParticles; // flags for MC collisions with HF particles
83+
static constexpr double mass{o2::constants::physics::MassBPlus};
9384

9485
using CollisionsWCentMult = soa::Join<aod::Collisions, aod::CentFV0As, aod::CentFT0Ms, aod::CentFT0As, aod::CentFT0Cs, aod::PVMultZeqs>;
9586
using CollisionsWMcCentMult = soa::Join<aod::Collisions, aod::McCollisionLabels, aod::CentFV0As, aod::CentFT0Ms, aod::CentFT0As, aod::CentFT0Cs, aod::PVMultZeqs>;
@@ -129,73 +120,14 @@ struct HfDerivedDataCreatorBplusToD0Pi {
129120
if (std::accumulate(doprocess.begin(), doprocess.end(), 0) != 1) {
130121
LOGP(fatal, "Only one process function can be enabled at a time.");
131122
}
132-
}
133-
134-
template <bool isMC, typename T>
135-
// void fillTablesCollision(const T& collision, int isEventReject, int runNumber)
136-
void fillTablesCollision(const T& collision)
137-
{
138-
if (fillCollBase) {
139-
rowCollBase(
140-
collision.posX(),
141-
collision.posY(),
142-
collision.posZ(),
143-
collision.numContrib(),
144-
collision.centFT0A(),
145-
collision.centFT0C(),
146-
collision.centFT0M(),
147-
collision.centFV0A(),
148-
collision.multZeqNTracksPV());
149-
// isEventReject,
150-
// runNumber);
151-
}
152-
if (fillCollId) {
153-
rowCollId(
154-
collision.globalIndex());
155-
}
156-
if constexpr (isMC) {
157-
if (fillMcRCollId && collision.has_mcCollision()) {
158-
// Save rowCollBase.lastIndex() at key collision.mcCollisionId()
159-
LOGF(debug, "Rec. collision %d: Filling derived-collision index %d for MC collision %d", collision.globalIndex(), rowCollBase.lastIndex(), collision.mcCollisionId());
160-
matchedCollisions[collision.mcCollisionId()].push_back(rowCollBase.lastIndex()); // [] inserts an empty element if it does not exist
161-
}
162-
}
163-
}
164-
165-
template <typename T>
166-
void fillTablesMcCollision(const T& mcCollision)
167-
{
168-
if (fillMcCollBase) {
169-
rowMcCollBase(
170-
mcCollision.posX(),
171-
mcCollision.posY(),
172-
mcCollision.posZ(),
173-
mcCollision.centFT0M());
174-
}
175-
if (fillMcCollId) {
176-
rowMcCollId(
177-
mcCollision.globalIndex());
178-
}
179-
if (fillMcRCollId) {
180-
// Fill the table with the vector of indices of derived reconstructed collisions matched to mcCollision.globalIndex()
181-
rowMcRCollId(
182-
matchedCollisions[mcCollision.globalIndex()]);
183-
}
123+
rowsCommon.init(confDerData);
184124
}
185125

186126
template <typename T, typename U, typename V>
187127
void fillTablesCandidate(const T& candidate, const U& prongCharm, const V& prongBachelor, int candFlag, double invMass,
188128
double ct, double y, int8_t flagMc, int8_t origin, float mlScore, const std::vector<float>& mlScoresCharm)
189129
{
190-
if (fillCandidateBase) {
191-
rowCandidateBase(
192-
rowCollBase.lastIndex(),
193-
candidate.pt(),
194-
candidate.eta(),
195-
candidate.phi(),
196-
invMass,
197-
y);
198-
}
130+
rowsCommon.fillTablesCandidate(candidate, invMass, y);
199131
if (fillCandidatePar) {
200132
rowCandidatePar(
201133
candidate.chi2PCA(),
@@ -294,26 +226,6 @@ struct HfDerivedDataCreatorBplusToD0Pi {
294226
}
295227
}
296228

297-
template <typename T, typename U>
298-
void fillTablesParticle(const T& particle, U mass)
299-
{
300-
if (fillParticleBase) {
301-
rowParticleBase(
302-
rowMcCollBase.lastIndex(),
303-
particle.pt(),
304-
particle.eta(),
305-
particle.phi(),
306-
RecoDecayPtEtaPhi::y(particle.pt(), particle.eta(), mass),
307-
particle.flagMcMatchGen(),
308-
particle.originMcGen());
309-
}
310-
if (fillParticleId) {
311-
rowParticleId(
312-
particle.mcCollisionId(),
313-
particle.globalIndex());
314-
}
315-
}
316-
317229
template <bool isMl, bool isMc, bool onlyBkg, bool onlySig, typename CollType, typename CandType, typename CandCharmType>
318230
void processCandidates(CollType const& collisions,
319231
Partition<CandType>& candidates,
@@ -323,13 +235,12 @@ struct HfDerivedDataCreatorBplusToD0Pi {
323235
{
324236
// Fill collision properties
325237
if constexpr (isMc) {
326-
if (fillMcRCollId) {
327-
matchedCollisions.clear();
238+
if (confDerData.fillMcRCollId) {
239+
rowsCommon.matchedCollisions.clear();
328240
}
329241
}
330242
auto sizeTableColl = collisions.size();
331-
reserveTable(rowCollBase, fillCollBase, sizeTableColl);
332-
reserveTable(rowCollId, fillCollId, sizeTableColl);
243+
rowsCommon.reserveTablesColl(sizeTableColl);
333244
for (const auto& collision : collisions) {
334245
auto thisCollId = collision.globalIndex();
335246
auto candidatesThisColl = candidates->sliceByCached(aod::hf_cand::collisionId, thisCollId, cache); // FIXME
@@ -338,19 +249,18 @@ struct HfDerivedDataCreatorBplusToD0Pi {
338249
// Skip collisions without HF candidates (and without HF particles in matched MC collisions if saving indices of reconstructed collisions matched to MC collisions)
339250
bool mcCollisionHasMcParticles{false};
340251
if constexpr (isMc) {
341-
mcCollisionHasMcParticles = fillMcRCollId && collision.has_mcCollision() && hasMcParticles[collision.mcCollisionId()];
252+
mcCollisionHasMcParticles = confDerData.fillMcRCollId && collision.has_mcCollision() && rowsCommon.hasMcParticles[collision.mcCollisionId()];
342253
LOGF(debug, "Rec. collision %d has MC collision %d with MC particles? %s", thisCollId, collision.mcCollisionId(), mcCollisionHasMcParticles ? "yes" : "no");
343254
}
344-
if (sizeTableCand == 0 && (!fillMcRCollId || !mcCollisionHasMcParticles)) {
255+
if (sizeTableCand == 0 && (!confDerData.fillMcRCollId || !mcCollisionHasMcParticles)) {
345256
LOGF(debug, "Skipping rec. collision %d", thisCollId);
346257
continue;
347258
}
348-
LOGF(debug, "Filling rec. collision %d at derived index %d", thisCollId, rowCollBase.lastIndex() + 1);
349-
// fillTablesCollision(collision, 0, collision.bc().runNumber());
350-
fillTablesCollision<isMc>(collision);
259+
LOGF(debug, "Filling rec. collision %d at derived index %d", thisCollId, rowsCommon.rowCollBase.lastIndex() + 1);
260+
rowsCommon.fillTablesCollision<isMc>(collision);
351261

352262
// Fill candidate properties
353-
reserveTable(rowCandidateBase, fillCandidateBase, sizeTableCand);
263+
rowsCommon.reserveTablesCandidates(sizeTableCand);
354264
reserveTable(rowCandidatePar, fillCandidatePar, sizeTableCand);
355265
reserveTable(rowCandidateParD0, fillCandidateParD0, sizeTableCand);
356266
reserveTable(rowCandidateParE, fillCandidateParE, sizeTableCand);
@@ -405,55 +315,6 @@ struct HfDerivedDataCreatorBplusToD0Pi {
405315
}
406316
}
407317

408-
template <typename CollisionType, typename ParticleType>
409-
void preProcessMcCollisions(CollisionType const& mcCollisions,
410-
ParticleType const& mcParticles)
411-
{
412-
if (!fillMcRCollId) {
413-
return;
414-
}
415-
hasMcParticles.clear();
416-
// Fill MC collision flags
417-
for (const auto& mcCollision : mcCollisions) {
418-
auto thisMcCollId = mcCollision.globalIndex();
419-
auto particlesThisMcColl = mcParticles.sliceBy(mcParticlesPerMcCollision, thisMcCollId);
420-
LOGF(debug, "MC collision %d has %d MC particles (preprocess)", thisMcCollId, particlesThisMcColl.size());
421-
hasMcParticles[thisMcCollId] = (particlesThisMcColl.size() > 0);
422-
}
423-
}
424-
425-
template <typename CollisionType, typename ParticleType>
426-
void processMcParticles(CollisionType const& mcCollisions,
427-
ParticleType const& mcParticles)
428-
{
429-
// Fill MC collision properties
430-
auto sizeTableMcColl = mcCollisions.size();
431-
reserveTable(rowMcCollBase, fillMcCollBase, sizeTableMcColl);
432-
reserveTable(rowMcCollId, fillMcCollId, sizeTableMcColl);
433-
reserveTable(rowMcRCollId, fillMcRCollId, sizeTableMcColl);
434-
for (const auto& mcCollision : mcCollisions) {
435-
auto thisMcCollId = mcCollision.globalIndex();
436-
auto particlesThisMcColl = mcParticles.sliceBy(mcParticlesPerMcCollision, thisMcCollId);
437-
auto sizeTablePart = particlesThisMcColl.size();
438-
LOGF(debug, "MC collision %d has %d MC particles", thisMcCollId, sizeTablePart);
439-
// Skip MC collisions without HF particles (and without HF candidates in matched reconstructed collisions if saving indices of reconstructed collisions matched to MC collisions)
440-
LOGF(debug, "MC collision %d has %d saved derived rec. collisions", thisMcCollId, matchedCollisions[thisMcCollId].size());
441-
if (sizeTablePart == 0 && (!fillMcRCollId || matchedCollisions[thisMcCollId].empty())) {
442-
LOGF(debug, "Skipping MC collision %d", thisMcCollId);
443-
continue;
444-
}
445-
LOGF(debug, "Filling MC collision %d at derived index %d", thisMcCollId, rowMcCollBase.lastIndex() + 1);
446-
fillTablesMcCollision(mcCollision);
447-
448-
// Fill MC particle properties
449-
reserveTable(rowParticleBase, fillParticleBase, sizeTablePart);
450-
reserveTable(rowParticleId, fillParticleId, sizeTablePart);
451-
for (const auto& particle : particlesThisMcColl) {
452-
fillTablesParticle(particle, o2::constants::physics::MassBPlus);
453-
}
454-
}
455-
}
456-
457318
void processData(CollisionsWCentMult const& collisions,
458319
SelectedCandidates const&,
459320
THfCandDaughters const& candidatesDaughters,
@@ -472,9 +333,9 @@ struct HfDerivedDataCreatorBplusToD0Pi {
472333
TracksWPid const& tracks,
473334
aod::BCs const& bcs)
474335
{
475-
preProcessMcCollisions(mcCollisions, mcParticles);
336+
rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
476337
processCandidates<false, true, false, true>(collisions, candidatesMcSig, candidatesDaughters, tracks, bcs);
477-
processMcParticles(mcCollisions, mcParticles);
338+
rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, mass);
478339
}
479340
PROCESS_SWITCH(HfDerivedDataCreatorBplusToD0Pi, processMcSig, "Process MC only for signals", false);
480341

@@ -486,9 +347,9 @@ struct HfDerivedDataCreatorBplusToD0Pi {
486347
TracksWPid const& tracks,
487348
aod::BCs const& bcs)
488349
{
489-
preProcessMcCollisions(mcCollisions, mcParticles);
350+
rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
490351
processCandidates<false, true, true, false>(collisions, candidatesMcBkg, candidatesDaughters, tracks, bcs);
491-
processMcParticles(mcCollisions, mcParticles);
352+
rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, mass);
492353
}
493354
PROCESS_SWITCH(HfDerivedDataCreatorBplusToD0Pi, processMcBkg, "Process MC only for background", false);
494355

@@ -500,9 +361,9 @@ struct HfDerivedDataCreatorBplusToD0Pi {
500361
TracksWPid const& tracks,
501362
aod::BCs const& bcs)
502363
{
503-
preProcessMcCollisions(mcCollisions, mcParticles);
364+
rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
504365
processCandidates<false, true, false, false>(collisions, candidatesMcAll, candidatesDaughters, tracks, bcs);
505-
processMcParticles(mcCollisions, mcParticles);
366+
rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, mass);
506367
}
507368
PROCESS_SWITCH(HfDerivedDataCreatorBplusToD0Pi, processMcAll, "Process MC", false);
508369

@@ -526,9 +387,9 @@ struct HfDerivedDataCreatorBplusToD0Pi {
526387
TracksWPid const& tracks,
527388
aod::BCs const& bcs)
528389
{
529-
preProcessMcCollisions(mcCollisions, mcParticles);
390+
rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
530391
processCandidates<true, true, false, true>(collisions, candidatesMcMlSig, candidatesDaughters, tracks, bcs);
531-
processMcParticles(mcCollisions, mcParticles);
392+
rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, mass);
532393
}
533394
PROCESS_SWITCH(HfDerivedDataCreatorBplusToD0Pi, processMcMlSig, "Process MC with ML only for signals", false);
534395

@@ -540,9 +401,9 @@ struct HfDerivedDataCreatorBplusToD0Pi {
540401
TracksWPid const& tracks,
541402
aod::BCs const& bcs)
542403
{
543-
preProcessMcCollisions(mcCollisions, mcParticles);
404+
rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
544405
processCandidates<true, true, true, false>(collisions, candidatesMcMlBkg, candidatesDaughters, tracks, bcs);
545-
processMcParticles(mcCollisions, mcParticles);
406+
rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, mass);
546407
}
547408
PROCESS_SWITCH(HfDerivedDataCreatorBplusToD0Pi, processMcMlBkg, "Process MC with ML only for background", false);
548409

@@ -554,9 +415,9 @@ struct HfDerivedDataCreatorBplusToD0Pi {
554415
TracksWPid const& tracks,
555416
aod::BCs const& bcs)
556417
{
557-
preProcessMcCollisions(mcCollisions, mcParticles);
418+
rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
558419
processCandidates<true, true, false, false>(collisions, candidatesMcMlAll, candidatesDaughters, tracks, bcs);
559-
processMcParticles(mcCollisions, mcParticles);
420+
rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, mass);
560421
}
561422
PROCESS_SWITCH(HfDerivedDataCreatorBplusToD0Pi, processMcMlAll, "Process MC with ML", false);
562423
};

0 commit comments

Comments
 (0)