Skip to content

Commit 4f3bb05

Browse files
skundu692skundu692
committed
Add optimized mixed event process function
1 parent 4f52188 commit 4f3bb05

File tree

1 file changed

+125
-30
lines changed

1 file changed

+125
-30
lines changed

‎PWGLF/Tasks/Strangeness/lambdaspincorrderived.cxx‎

Lines changed: 125 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
#include "Framework/ASoAHelpers.h"
2323
#include "Framework/AnalysisDataModel.h"
2424
#include "Framework/AnalysisTask.h"
25+
#include "Framework/BinningPolicy.h"
2526
#include "Framework/StepTHn.h"
2627
#include "Framework/runDataProcessing.h"
2728
#include <Framework/Configurable.h>
@@ -34,9 +35,14 @@
3435

3536
#include <fairlogger/Logger.h>
3637

38+
#include <cmath> // for std::fabs
39+
#include <deque>
3740
#include <iostream>
3841
#include <iterator>
42+
#include <set> // <<< CHANGED: for dedup sets
3943
#include <string>
44+
#include <type_traits>
45+
#include <unordered_map> // <<< CHANGED: for seenMap
4046
#include <vector>
4147

4248
// o2 includes.
@@ -49,7 +55,7 @@ using namespace o2::framework::expressions;
4955
using namespace o2::soa;
5056

5157
struct lambdaspincorrderived {
52-
58+
// BinningType colBinning;
5359
struct : ConfigurableGroup {
5460
Configurable<std::string> cfgURL{"cfgURL", "http://alice-ccdb.cern.ch", "Address of the CCDB to browse"};
5561
Configurable<int64_t> nolaterthan{"ccdb-no-later-than", std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count(), "Latest acceptable timestamp of creation for the object"};
@@ -71,6 +77,7 @@ struct lambdaspincorrderived {
7177
Configurable<float> centMax{"centMax", 80, "Maximum Centrality"};
7278

7379
// Lambda selection ////////////
80+
Configurable<unsigned> harmonic{"harmonic", 1, "Harmonic delta phi"};
7481
Configurable<bool> useweight{"useweight", 1, "Use weight"};
7582
Configurable<bool> usePDGM{"usePDGM", 1, "Use PDG mass"};
7683
Configurable<bool> checkDoubleStatus{"checkDoubleStatus", 0, "Check Double status"};
@@ -183,7 +190,7 @@ struct lambdaspincorrderived {
183190
if (std::abs(candidate1.lambdaEta() - candidate2.lambdaEta()) > etaMix) {
184191
return false;
185192
}
186-
if (std::abs(RecoDecay::constrainAngle(candidate1.lambdaPhi(), 0.0F) - RecoDecay::constrainAngle(candidate2.lambdaPhi(), 0.0F)) > phiMix) {
193+
if (std::abs(RecoDecay::constrainAngle(candidate1.lambdaPhi(), 0.0F, harmonic) - RecoDecay::constrainAngle(candidate2.lambdaPhi(), 0.0F, harmonic)) > phiMix) {
187194
return false;
188195
}
189196
if (std::abs(candidate1.lambdaMass() - candidate2.lambdaMass()) > massMix) {
@@ -195,7 +202,7 @@ struct lambdaspincorrderived {
195202
void fillHistograms(int tag1, int tag2,
196203
const ROOT::Math::PtEtaPhiMVector& particle1, const ROOT::Math::PtEtaPhiMVector& particle2,
197204
const ROOT::Math::PtEtaPhiMVector& daughpart1, const ROOT::Math::PtEtaPhiMVector& daughpart2,
198-
double centrality, int datatype)
205+
double centrality, int datatype, float mixpairweight)
199206
{
200207

201208
auto lambda1Mass = 0.0;
@@ -230,41 +237,42 @@ struct lambdaspincorrderived {
230237

231238
auto cosThetaDiff = -999.0;
232239
cosThetaDiff = proton1LambdaRF.Vect().Unit().Dot(proton2LambdaRF.Vect().Unit());
233-
double deltaPhi = std::abs(RecoDecay::constrainAngle(particle1Dummy.Phi(), 0.0F) - RecoDecay::constrainAngle(particle2Dummy.Phi(), 0.0F));
240+
double deltaPhi = std::abs(RecoDecay::constrainAngle(particle1Dummy.Phi(), 0.0F, harmonic) - RecoDecay::constrainAngle(particle2Dummy.Phi(), 0.0F, harmonic));
234241
double deltaEta = particle1Dummy.Eta() - particle2Dummy.Eta();
235242
double deltaR = TMath::Sqrt(deltaEta * deltaEta + deltaPhi * deltaPhi);
236243

237244
if (datatype == 0) {
238-
histos.fill(HIST("hPtYSame"), particle1.Pt(), particle1.Rapidity());
245+
mixpairweight = 1.0;
246+
histos.fill(HIST("hPtYSame"), particle1.Pt(), particle1.Rapidity(), mixpairweight);
239247
if (tag1 == 0 && tag2 == 0) {
240-
histos.fill(HIST("hSparseLambdaLambda"), particle1.M(), particle2.M(), cosThetaDiff, centrality, deltaR);
241-
histos.fill(HIST("hLambdaSameForLL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F));
248+
histos.fill(HIST("hSparseLambdaLambda"), particle1.M(), particle2.M(), cosThetaDiff, centrality, deltaR, mixpairweight);
249+
histos.fill(HIST("hLambdaSameForLL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F, harmonic), mixpairweight);
242250
} else if ((tag1 == 0 && tag2 == 1) || (tag1 == 1 && tag2 == 0)) {
243-
histos.fill(HIST("hSparseLambdaAntiLambda"), particle1.M(), particle2.M(), cosThetaDiff, centrality, deltaR);
244-
histos.fill(HIST("hLambdaSameForLAL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F));
251+
histos.fill(HIST("hSparseLambdaAntiLambda"), particle1.M(), particle2.M(), cosThetaDiff, centrality, deltaR, mixpairweight);
252+
histos.fill(HIST("hLambdaSameForLAL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F, harmonic), mixpairweight);
245253
} else if (tag1 == 1 && tag2 == 1) {
246-
histos.fill(HIST("hSparseAntiLambdaAntiLambda"), particle1.M(), particle2.M(), cosThetaDiff, centrality, deltaR);
247-
histos.fill(HIST("hAntiLambdaSameForALAL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F));
254+
histos.fill(HIST("hSparseAntiLambdaAntiLambda"), particle1.M(), particle2.M(), cosThetaDiff, centrality, deltaR, mixpairweight);
255+
histos.fill(HIST("hAntiLambdaSameForALAL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F, harmonic), mixpairweight);
248256
}
249257
} else if (datatype == 1) {
250-
double weight1 = 1.0;
251-
double weight2 = 1.0;
252-
double weight3 = 1.0;
258+
double weight1 = mixpairweight;
259+
double weight2 = mixpairweight;
260+
double weight3 = mixpairweight;
253261
if (useweight) {
254-
weight1 = hweight1->GetBinContent(hweight1->FindBin(particle1.Pt(), particle1.Eta(), particle1.Phi()));
255-
weight2 = hweight2->GetBinContent(hweight2->FindBin(particle1.Pt(), particle1.Eta(), particle1.Phi()));
256-
weight3 = hweight3->GetBinContent(hweight3->FindBin(particle1.Pt(), particle1.Eta(), particle1.Phi()));
262+
weight1 = mixpairweight * hweight1->GetBinContent(hweight1->FindBin(particle1.Pt(), particle1.Eta(), particle1.Phi()));
263+
weight2 = mixpairweight * hweight2->GetBinContent(hweight2->FindBin(particle1.Pt(), particle1.Eta(), particle1.Phi()));
264+
weight3 = mixpairweight * hweight3->GetBinContent(hweight3->FindBin(particle1.Pt(), particle1.Eta(), particle1.Phi()));
257265
}
258266
histos.fill(HIST("hPtYMix"), particle1.Pt(), particle1.Rapidity());
259267
if (tag1 == 0 && tag2 == 0) {
260268
histos.fill(HIST("hSparseLambdaLambdaMixed"), particle1.M(), particle2.M(), cosThetaDiff, centrality, deltaR, weight1);
261-
histos.fill(HIST("hLambdaMixForLL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F), weight1);
269+
histos.fill(HIST("hLambdaMixForLL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F, harmonic), weight1);
262270
} else if ((tag1 == 0 && tag2 == 1) || (tag1 == 1 && tag2 == 0)) {
263271
histos.fill(HIST("hSparseLambdaAntiLambdaMixed"), particle1.M(), particle2.M(), cosThetaDiff, centrality, deltaR, weight2);
264-
histos.fill(HIST("hLambdaMixForLAL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F), weight2);
272+
histos.fill(HIST("hLambdaMixForLAL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F, harmonic), weight2);
265273
} else if (tag1 == 1 && tag2 == 1) {
266274
histos.fill(HIST("hSparseAntiLambdaAntiLambdaMixed"), particle1.M(), particle2.M(), cosThetaDiff, centrality, deltaR, weight3);
267-
histos.fill(HIST("hAntiLambdaMixForALAL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F), weight3);
275+
histos.fill(HIST("hAntiLambdaMixForALAL"), particle1.Pt(), particle1.Eta(), RecoDecay::constrainAngle(particle1.Phi(), 0.0F, harmonic), weight3);
268276
}
269277
}
270278
}
@@ -310,24 +318,24 @@ struct lambdaspincorrderived {
310318
}
311319
proton2 = ROOT::Math::PtEtaPhiMVector(v02.protonPt(), v02.protonEta(), v02.protonPhi(), o2::constants::physics::MassProton);
312320
lambda2 = ROOT::Math::PtEtaPhiMVector(v02.lambdaPt(), v02.lambdaEta(), v02.lambdaPhi(), v02.lambdaMass());
313-
histos.fill(HIST("deltaPhiSame"), std::abs(RecoDecay::constrainAngle(v0.lambdaPhi(), 0.0F) - RecoDecay::constrainAngle(v02.lambdaPhi(), 0.0F)));
321+
histos.fill(HIST("deltaPhiSame"), std::abs(RecoDecay::constrainAngle(v0.lambdaPhi(), 0.0F, harmonic) - RecoDecay::constrainAngle(v02.lambdaPhi(), 0.0F, harmonic)));
314322
if (std::abs(lambda2.Rapidity()) > rapidity) {
315323
continue;
316324
}
317325
if (std::abs(lambda2.Eta()) > v0eta) {
318326
continue;
319327
}
320328
if (v0.v0Status() == 0 && v02.v0Status() == 0) {
321-
fillHistograms(0, 0, lambda, lambda2, proton, proton2, centrality, 0);
329+
fillHistograms(0, 0, lambda, lambda2, proton, proton2, centrality, 0, 1.0);
322330
}
323331
if (v0.v0Status() == 0 && v02.v0Status() == 1) {
324-
fillHistograms(0, 1, lambda, lambda2, proton, proton2, centrality, 0);
332+
fillHistograms(0, 1, lambda, lambda2, proton, proton2, centrality, 0, 1.0);
325333
}
326334
if (v0.v0Status() == 1 && v02.v0Status() == 0) {
327-
fillHistograms(1, 0, lambda2, lambda, proton2, proton, centrality, 0);
335+
fillHistograms(1, 0, lambda2, lambda, proton2, proton, centrality, 0, 1.0);
328336
}
329337
if (v0.v0Status() == 1 && v02.v0Status() == 1) {
330-
fillHistograms(1, 1, lambda, lambda2, proton, proton2, centrality, 0);
338+
fillHistograms(1, 1, lambda, lambda2, proton, proton2, centrality, 0, 1.0);
331339
}
332340
}
333341
}
@@ -388,7 +396,7 @@ struct lambdaspincorrderived {
388396
lambda = ROOT::Math::PtEtaPhiMVector(t3.lambdaPt(), t3.lambdaEta(), t3.lambdaPhi(), t3.lambdaMass());
389397
proton2 = ROOT::Math::PtEtaPhiMVector(t2.protonPt(), t2.protonEta(), t2.protonPhi(), o2::constants::physics::MassProton);
390398
lambda2 = ROOT::Math::PtEtaPhiMVector(t2.lambdaPt(), t2.lambdaEta(), t2.lambdaPhi(), t2.lambdaMass());
391-
histos.fill(HIST("deltaPhiMix"), std::abs(RecoDecay::constrainAngle(t3.lambdaPhi(), 0.0F) - RecoDecay::constrainAngle(t2.lambdaPhi(), 0.0F)));
399+
histos.fill(HIST("deltaPhiMix"), std::abs(RecoDecay::constrainAngle(t3.lambdaPhi(), 0.0F, harmonic) - RecoDecay::constrainAngle(t2.lambdaPhi(), 0.0F, harmonic)));
392400
if (std::abs(lambda.Rapidity()) > rapidity) {
393401
continue;
394402
}
@@ -402,22 +410,109 @@ struct lambdaspincorrderived {
402410
continue;
403411
}
404412
if (t3.v0Status() == 0 && t2.v0Status() == 0) {
405-
fillHistograms(0, 0, lambda, lambda2, proton, proton2, centrality, 1);
413+
fillHistograms(0, 0, lambda, lambda2, proton, proton2, centrality, 1, 1.0);
406414
}
407415
if (t3.v0Status() == 0 && t2.v0Status() == 1) {
408-
fillHistograms(0, 1, lambda, lambda2, proton, proton2, centrality, 1);
416+
fillHistograms(0, 1, lambda, lambda2, proton, proton2, centrality, 1, 1.0);
409417
}
410418
if (t3.v0Status() == 1 && t2.v0Status() == 0) {
411-
fillHistograms(1, 0, lambda2, lambda, proton2, proton, centrality, 1);
419+
fillHistograms(1, 0, lambda2, lambda, proton2, proton, centrality, 1, 1.0);
412420
}
413421
if (t3.v0Status() == 1 && t2.v0Status() == 1) {
414-
fillHistograms(1, 1, lambda, lambda2, proton, proton2, centrality, 1);
422+
fillHistograms(1, 1, lambda, lambda2, proton, proton2, centrality, 1, 1.0);
415423
}
416424
}
417425
} // replacement track pair
418426
} // collision pair
419427
}
420428
PROCESS_SWITCH(lambdaspincorrderived, processME, "Process data ME", false);
429+
430+
void processMEV2(EventCandidates const& collisions, AllTrackCandidates const& V0s)
431+
{
432+
auto nBins = colBinning.getAllBinsCount();
433+
std::vector<std::deque<std::pair<int, AllTrackCandidates>>> eventPools(nBins);
434+
435+
for (auto& collision1 : collisions) {
436+
int bin = colBinning.getBin(std::make_tuple(collision1.posz(), collision1.cent()));
437+
auto poolA = V0s.sliceBy(tracksPerCollisionV0, collision1.index());
438+
float centrality = collision1.cent();
439+
440+
// <<< CHANGED: map old collision index → set of (t2.idx, t3.idx) we've already filled
441+
std::unordered_map<int, std::set<std::pair<int, int>>> seenMap;
442+
443+
for (auto& [t1, t2] : soa::combinations(o2::soa::CombinationsFullIndexPolicy(poolA, poolA))) {
444+
if (!selectionV0(t1) || !selectionV0(t2))
445+
continue;
446+
if (t2.index() <= t1.index())
447+
continue;
448+
if (t1.protonIndex() == t2.protonIndex())
449+
continue;
450+
if (t1.pionIndex() == t2.pionIndex())
451+
continue;
452+
453+
int mixes = 0;
454+
for (auto it = eventPools[bin].rbegin(); it != eventPools[bin].rend() && mixes < nEvtMixing; ++it, ++mixes) {
455+
int collision2idx = it->first;
456+
AllTrackCandidates& poolB = it->second;
457+
458+
int nRepl = 0;
459+
for (auto& t3 : poolB) {
460+
if (selectionV0(t3) && checkKinematics(t1, t3)) {
461+
++nRepl;
462+
}
463+
}
464+
if (nRepl == 0)
465+
continue;
466+
float invN = 1.0f / float(nRepl);
467+
468+
for (auto& t3 : poolB) {
469+
if (!(selectionV0(t3) && checkKinematics(t1, t3))) {
470+
continue;
471+
}
472+
if (collision1.index() == collision2idx) {
473+
continue;
474+
}
475+
476+
// <<< CHANGED: dedupe (t2, t3) pairs per prior collision
477+
auto key = std::make_pair(t2.index(), t3.index());
478+
auto& seen = seenMap[collision2idx];
479+
if (!seen.insert(key).second) {
480+
continue;
481+
}
482+
483+
// reconstruct 4-vectors
484+
auto proton = ROOT::Math::PtEtaPhiMVector(t3.protonPt(), t3.protonEta(), t3.protonPhi(), o2::constants::physics::MassProton);
485+
auto lambda = ROOT::Math::PtEtaPhiMVector(t3.lambdaPt(), t3.lambdaEta(), t3.lambdaPhi(), t3.lambdaMass());
486+
auto proton2 = ROOT::Math::PtEtaPhiMVector(t2.protonPt(), t2.protonEta(), t2.protonPhi(), o2::constants::physics::MassProton);
487+
auto lambda2 = ROOT::Math::PtEtaPhiMVector(t2.lambdaPt(), t2.lambdaEta(), t2.lambdaPhi(), t2.lambdaMass());
488+
489+
float dPhi = std::fabs(RecoDecay::constrainAngle(lambda.Phi(), 0.0F, harmonic) - RecoDecay::constrainAngle(lambda2.Phi(), 0.0F, harmonic));
490+
histos.fill(HIST("deltaPhiMix"), dPhi, invN);
491+
492+
if (t3.v0Status() == 0 && t2.v0Status() == 0) {
493+
fillHistograms(0, 0, lambda, lambda2, proton, proton2, centrality, 1, invN);
494+
}
495+
if (t3.v0Status() == 0 && t2.v0Status() == 1) {
496+
fillHistograms(0, 1, lambda, lambda2, proton, proton2, centrality, 1, invN);
497+
}
498+
if (t3.v0Status() == 1 && t2.v0Status() == 0) {
499+
fillHistograms(1, 0, lambda2, lambda, proton2, proton, centrality, 1, invN);
500+
}
501+
if (t3.v0Status() == 1 && t2.v0Status() == 1) {
502+
fillHistograms(1, 1, lambda, lambda2, proton, proton2, centrality, 1, invN);
503+
}
504+
}
505+
} // end mixing-event loop
506+
} // end same-event pair loop
507+
508+
auto sliced = V0s.sliceBy(tracksPerCollisionV0, collision1.index());
509+
eventPools[bin].emplace_back(collision1.index(), std::move(sliced));
510+
if ((int)eventPools[bin].size() > nEvtMixing) {
511+
eventPools[bin].pop_front();
512+
}
513+
} // end primary-event loop
514+
}
515+
PROCESS_SWITCH(lambdaspincorrderived, processMEV2, "Process data ME", false);
421516
};
422517
WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
423518
{

0 commit comments

Comments
 (0)