Skip to content

Commit 518dc5c

Browse files
zjxiongOvOzjxiongOvO
committed
speed up runMCGen
1 parent fced339 commit 518dc5c

File tree

1 file changed

+61
-50
lines changed

1 file changed

+61
-50
lines changed

‎PWGDQ/Tasks/dqEfficiency_withAssoc.cxx‎

Lines changed: 61 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1370,6 +1370,7 @@ struct AnalysisSameEventPairing {
13701370
Configurable<std::string> genSignalsJSON{"cfgMCGenSignalsJSON", "", "Additional list of MC signals (generated) via JSON"};
13711371
Configurable<std::string> recSignals{"cfgBarrelMCRecSignals", "", "Comma separated list of MC signals (reconstructed)"};
13721372
Configurable<std::string> recSignalsJSON{"cfgMCRecSignalsJSON", "", "Comma separated list of MC signals (reconstructed) via JSON"};
1373+
Configurable<std::string> finalStateSignals{"cfgBarrelMCFinalStateSignals", "eFromJpsi", "Comma separated list of MC signals (final state particles)"};
13731374
Configurable<bool> skimSignalOnly{"cfgSkimSignalOnly", false, "Configurable to select only matched candidates"};
13741375
} fConfigMC;
13751376

@@ -1396,6 +1397,7 @@ struct AnalysisSameEventPairing {
13961397
std::map<int, std::vector<TString>> fMuonHistNamesMCmatched;
13971398
std::vector<MCSignal*> fRecMCSignals;
13981399
std::vector<MCSignal*> fGenMCSignals;
1400+
std::vector<MCSignal*> fFinalStateMCSignals;
13991401

14001402
std::vector<AnalysisCompositeCut> fPairCuts;
14011403
std::vector<AnalysisCut*> fMCGenAccCuts;
@@ -1691,6 +1693,16 @@ struct AnalysisSameEventPairing {
16911693
}
16921694
}
16931695

1696+
// define final state MC signals
1697+
TString sigFinalStateNamesStr = fConfigMC.finalStateSignals.value;
1698+
std::unique_ptr<TObjArray> objFinalStateSigArray(sigFinalStateNamesStr.Tokenize(","));
1699+
for (int isig = 0; isig < objFinalStateSigArray->GetEntries(); isig++) {
1700+
MCSignal* sig = o2::aod::dqmcsignals::GetMCSignal(objFinalStateSigArray->At(isig)->GetName());
1701+
if (sig) {
1702+
fFinalStateMCSignals.push_back(sig);
1703+
}
1704+
}
1705+
16941706
if (isMCGen) {
16951707
for (auto& sig : fGenMCSignals) {
16961708
if (sig->GetNProngs() == 1) {
@@ -2161,23 +2173,27 @@ struct AnalysisSameEventPairing {
21612173
PresliceUnsorted<ReducedMCTracks> perReducedMcEvent = aod::reducedtrackMC::reducedMCeventId;
21622174

21632175
template <int TPairType>
2164-
void runMCGenWithGrouping(MyEventsVtxCovSelected const& events, ReducedMCEvents const& mcEvents, ReducedMCTracks const& mcTracks)
2176+
void runMCGen(MyEventsVtxCovSelected const& events, ReducedMCEvents const& mcEvents, ReducedMCTracks const& mcTracks)
21652177
{
2178+
cout << "AnalysisSameEventPairing::runMCGen() called" << endl;
21662179
uint32_t mcDecision = 0;
21672180
int isig = 0;
21682181

2182+
// Loop over all MC single particles to fill generator level histograms, disregarding of whether they belong to selected reconstructed events or not
21692183
for (auto& mctrack : mcTracks) {
2170-
VarManager::FillTrackMC(mcTracks, mctrack);
2171-
// NOTE: Signals are checked here mostly based on the skimmed MC stack, so depending on the requested signal, the stack could be incomplete.
2172-
// NOTE: However, the working model is that the decisions on MC signals are precomputed during skimming and are stored in the mcReducedFlags member.
2173-
// TODO: Use the mcReducedFlags to select signals
21742184
for (auto& sig : fGenMCSignals) {
21752185
if (sig->CheckSignal(true, mctrack)) {
2186+
VarManager::FillTrackMC(mcTracks, mctrack);
21762187
fHistMan->FillHistClass(Form("MCTruthGen_%s", sig->GetName()), VarManager::fgValues);
21772188
}
21782189
}
21792190
}
2180-
// Fill Generated histograms taking into account selected collisions
2191+
// cout << "Filled single MC particle generator histograms." << endl;
2192+
2193+
// make a vector of global indices of mc particles which fulfill the selected track-level signal definition (to speed up the two mc particle combinatorics)
2194+
std::vector<uint64_t> FinalStateMcParticleIndices;
2195+
2196+
// Now loop over reconstructed events to select only MC particles belonging to the same MC collision as the reconstructed event
21812197
for (auto& event : events) {
21822198
if (!event.isEventSelected_bit(0)) {
21832199
continue;
@@ -2186,16 +2202,21 @@ struct AnalysisSameEventPairing {
21862202
continue;
21872203
}
21882204

2189-
for (auto& track : mcTracks) {
2190-
if (track.reducedMCeventId() != event.reducedMCeventId()) {
2191-
continue;
2192-
}
2193-
VarManager::FillTrackMC(mcTracks, track);
2205+
FinalStateMcParticleIndices.clear();
2206+
2207+
auto groupedMCTracks = mcTracks.sliceBy(perReducedMcEvent, event.reducedMCeventId());
2208+
groupedMCTracks.bindInternalIndicesTo(&mcTracks);
2209+
2210+
for (auto& track : groupedMCTracks) {
21942211
auto track_raw = mcTracks.rawIteratorAt(track.globalIndex());
21952212
mcDecision = 0;
21962213
isig = 0;
21972214
for (auto& sig : fGenMCSignals) {
21982215
if (sig->CheckSignal(true, track_raw)) {
2216+
if (track.reducedMCeventId() != event.reducedMCeventId()) { // check that the mc track belongs to the same mc collision as the reconstructed event
2217+
continue;
2218+
}
2219+
VarManager::FillTrackMC(mcTracks, track);
21992220
mcDecision |= (static_cast<uint32_t>(1) << isig);
22002221
fHistMan->FillHistClass(Form("MCTruthGenSel_%s", sig->GetName()), VarManager::fgValues);
22012222
MCTruthTableEffi(VarManager::fgValues[VarManager::kMCPt], VarManager::fgValues[VarManager::kMCEta], VarManager::fgValues[VarManager::kMCY], VarManager::fgValues[VarManager::kMCPhi], VarManager::fgValues[VarManager::kMCVz], VarManager::fgValues[VarManager::kMCVtxZ], VarManager::fgValues[VarManager::kMultFT0A], VarManager::fgValues[VarManager::kMultFT0C], VarManager::fgValues[VarManager::kCentFT0M], VarManager::fgValues[VarManager::kVtxNcontribReal]);
@@ -2207,50 +2228,39 @@ struct AnalysisSameEventPairing {
22072228
}
22082229
isig++;
22092230
}
2210-
}
2211-
} // end loop over reconstructed events
2212-
2213-
if (fHasTwoProngGenMCsignals) {
2214-
for (auto& [t1, t2] : combinations(mcTracks, mcTracks)) {
2215-
auto t1_raw = mcTracks.rawIteratorAt(t1.globalIndex());
2216-
auto t2_raw = mcTracks.rawIteratorAt(t2.globalIndex());
2217-
if (t1_raw.reducedMCeventId() == t2_raw.reducedMCeventId()) {
2218-
for (auto& sig : fGenMCSignals) {
2219-
if (sig->GetNProngs() != 2) { // NOTE: 2-prong signals required here
2220-
continue;
2221-
}
2222-
if (sig->CheckSignal(true, t1_raw, t2_raw)) {
2223-
VarManager::FillPairMC<TPairType>(t1, t2);
2224-
fHistMan->FillHistClass(Form("MCTruthGenPair_%s", sig->GetName()), VarManager::fgValues);
2225-
}
2231+
for (auto& sig : fFinalStateMCSignals) {
2232+
if (sig->CheckSignal(true, track_raw)) {
2233+
FinalStateMcParticleIndices.push_back(track.globalIndex());
2234+
break; // if one of the final state signals is matched, no need to check the others
22262235
}
22272236
}
2237+
// cout << "Filled single MC particle generator histograms for reconstructed event." << endl;
22282238
}
2229-
}
2230-
for (auto& event : events) {
2231-
if (!event.isEventSelected_bit(0)) {
2232-
continue;
2233-
}
2234-
if (!event.has_reducedMCevent()) {
2235-
continue;
2236-
}
2237-
// CURRENTLY ONLY FOR 1-GENERATION 2-PRONG SIGNALS
2239+
22382240
if (fHasTwoProngGenMCsignals) {
2239-
auto groupedMCTracks = mcTracks.sliceBy(perReducedMcEvent, event.reducedMCeventId());
2240-
groupedMCTracks.bindInternalIndicesTo(&mcTracks);
2241-
for (auto& [t1, t2] : combinations(groupedMCTracks, groupedMCTracks)) {
2242-
auto t1_raw = mcTracks.rawIteratorAt(t1.globalIndex());
2243-
auto t2_raw = mcTracks.rawIteratorAt(t2.globalIndex());
2244-
if (t1_raw.reducedMCeventId() == t2_raw.reducedMCeventId()) {
2241+
// loop over combinations of the selected mc particles to fill generator level pair histograms
2242+
for (auto& t1 : FinalStateMcParticleIndices) {
2243+
auto t1_raw = mcTracks.rawIteratorAt(t1);
2244+
for (auto& t2 : FinalStateMcParticleIndices) {
2245+
if (t2 <= t1) {
2246+
continue; // avoid double counting and self-pairing
2247+
}
2248+
// for (auto& [t1, t2] : combinations(groupedMCTracks, groupedMCTracks)) {
2249+
// cout << "Processing pair of mcTracks with globalIndices = " << t1.globalIndex() << ", " << t2.globalIndex() << endl;
2250+
auto t2_raw = mcTracks.rawIteratorAt(t2);
2251+
22452252
mcDecision = 0;
22462253
isig = 0;
22472254
for (auto& sig : fGenMCSignals) {
22482255
if (sig->GetNProngs() != 2) { // NOTE: 2-prong signals required here
22492256
continue;
22502257
}
2258+
// cout << " Checking signal: " << sig->GetName() << endl;
22512259
if (sig->CheckSignal(true, t1_raw, t2_raw)) {
2260+
// cout << " Signal matched!" << endl;
22522261
mcDecision |= (static_cast<uint32_t>(1) << isig);
2253-
VarManager::FillPairMC<TPairType>(t1, t2);
2262+
VarManager::FillPairMC<TPairType>(t1_raw, t2_raw);
2263+
// cout << " Filled VarManager for the pair." << endl;
22542264
fHistMan->FillHistClass(Form("MCTruthGenPairSel_%s", sig->GetName()), VarManager::fgValues);
22552265
// Fill also acceptance cut histograms if requested
22562266
if (fUseMCGenAccCut) {
@@ -2262,15 +2272,16 @@ struct AnalysisSameEventPairing {
22622272
}
22632273
if (useMiniTree.fConfigMiniTree) {
22642274
// WARNING! To be checked
2265-
dileptonMiniTreeGen(mcDecision, -999, t1.pt(), t1.eta(), t1.phi(), t2.pt(), t2.eta(), t2.phi());
2275+
dileptonMiniTreeGen(mcDecision, -999, t1_raw.pt(), t1_raw.eta(), t1_raw.phi(), t2_raw.pt(), t2_raw.eta(), t2_raw.phi());
22662276
}
22672277
}
22682278
isig++;
22692279
}
2270-
}
2271-
}
2272-
} // end loop over reconstructed events
2273-
}
2280+
} // end loop over second mc particle
2281+
} // end loop over first mc particle
2282+
}
2283+
} // end loop over reconstructed events
2284+
// cout << "AnalysisSameEventPairing::runMCGen() completed" << endl;
22742285
}
22752286

22762287
void processAllSkimmed(MyEventsVtxCovSelected const& events,
@@ -2292,15 +2303,15 @@ struct AnalysisSameEventPairing {
22922303
MyBarrelTracksWithCovWithAmbiguities const& barrelTracks, ReducedMCEvents const& mcEvents, ReducedMCTracks const& mcTracks)
22932304
{
22942305
runSameEventPairing<true, VarManager::kDecayToEE, gkEventFillMapWithCov, gkTrackFillMapWithCov>(events, trackAssocsPerCollision, barrelAssocs, barrelTracks, mcEvents, mcTracks);
2295-
runMCGenWithGrouping<VarManager::kDecayToEE>(events, mcEvents, mcTracks);
2306+
runMCGen<VarManager::kDecayToEE>(events, mcEvents, mcTracks);
22962307
}
22972308

22982309
void processBarrelOnlyWithCollSkimmed(MyEventsVtxCovSelected const& events,
22992310
soa::Join<aod::ReducedTracksAssoc, aod::BarrelTrackCuts, aod::Prefilter> const& barrelAssocs,
23002311
MyBarrelTracksWithCovWithAmbiguitiesWithColl const& barrelTracks, ReducedMCEvents const& mcEvents, ReducedMCTracks const& mcTracks)
23012312
{
23022313
runSameEventPairing<true, VarManager::kDecayToEE, gkEventFillMapWithCov, gkTrackFillMapWithCovWithColl>(events, trackAssocsPerCollision, barrelAssocs, barrelTracks, mcEvents, mcTracks);
2303-
runMCGenWithGrouping<VarManager::kDecayToEE>(events, mcEvents, mcTracks);
2314+
runMCGen<VarManager::kDecayToEE>(events, mcEvents, mcTracks);
23042315
}
23052316

23062317
void processMuonOnlySkimmed(MyEventsVtxCovSelected const& events,

0 commit comments

Comments
 (0)