Skip to content

Commit 722aa42

Browse files
jaelparkjaelpark
authored
[PWGCF] Reduce correlations code duplication (#10205)
1 parent 23f3c30 commit 722aa42

File tree

2 files changed

+53
-106
lines changed

2 files changed

+53
-106
lines changed

‎PWGCF/TableProducer/filter2Prong.cxx‎

Lines changed: 36 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
// In applying this license CERN does not waive the privileges and immunities
99
// granted to it by virtue of its status as an Intergovernmental Organization
1010
// or submit itself to any jurisdiction.
11+
#include <experimental/type_traits>
1112
#include <vector>
1213

1314
#include "Framework/runDataProcessing.h"
@@ -39,10 +40,17 @@ struct Filter2Prong {
3940
Produces<aod::CF2ProngTracks> output2ProngTracks;
4041
Produces<aod::CF2ProngTrackmls> output2ProngTrackmls;
4142

43+
std::vector<float> mlvecd{};
44+
std::vector<float> mlvecdbar{};
45+
4246
using HFCandidates = soa::Join<aod::HfCand2Prong, aod::HfSelD0>;
4347
using HFCandidatesML = soa::Join<aod::HfCand2Prong, aod::HfSelD0, aod::HfMlD0>;
4448

45-
void processDataML(aod::Collisions::iterator const&, aod::BCsWithTimestamps const&, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidatesML const& candidates)
49+
template <class T>
50+
using HasMLProb = decltype(std::declval<T&>().mlProbD0());
51+
52+
template <class HFCandidatesType>
53+
void processDataT(aod::Collisions::iterator const&, aod::BCsWithTimestamps const&, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidatesType const& candidates)
4654
{
4755
if (cfcollisions.size() <= 0 || cftracks.size() <= 0)
4856
return; // rejected collision
@@ -68,71 +76,49 @@ struct Filter2Prong {
6876
if (cfgYMax >= 0.0f && std::abs(hfHelper.yD0(c)) > cfgYMax)
6977
continue;
7078

71-
std::vector<float> mlvecd{};
72-
std::vector<float> mlvecdbar{};
73-
7479
if (c.isSelD0() > 0) {
7580
output2ProngTracks(cfcollisions.begin().globalIndex(),
7681
prongCFId[0], prongCFId[1], c.pt(), c.eta(), c.phi(), hfHelper.invMassD0ToPiK(c), aod::cf2prongtrack::D0ToPiK);
77-
for (float val : c.mlProbD0()) {
78-
mlvecd.push_back(val);
82+
if constexpr (std::experimental::is_detected<HasMLProb, typename HFCandidatesType::iterator>::value) {
83+
mlvecd.clear();
84+
for (float val : c.mlProbD0()) {
85+
mlvecd.push_back(val);
86+
}
87+
mlvecdbar.clear();
88+
for (float val : c.mlProbD0bar()) {
89+
mlvecdbar.push_back(val);
90+
}
91+
output2ProngTrackmls(cfcollisions.begin().globalIndex(), mlvecd, mlvecdbar);
7992
}
80-
for (float val : c.mlProbD0bar()) {
81-
mlvecdbar.push_back(val);
82-
}
83-
output2ProngTrackmls(cfcollisions.begin().globalIndex(), mlvecd, mlvecdbar);
8493
}
8594

86-
mlvecd.clear();
87-
mlvecdbar.clear();
88-
8995
if (c.isSelD0bar() > 0) {
9096
output2ProngTracks(cfcollisions.begin().globalIndex(),
9197
prongCFId[0], prongCFId[1], c.pt(), c.eta(), c.phi(), hfHelper.invMassD0barToKPi(c), aod::cf2prongtrack::D0barToKPi);
92-
for (float val : c.mlProbD0()) {
93-
mlvecd.push_back(val);
98+
if constexpr (std::experimental::is_detected<HasMLProb, typename HFCandidatesType::iterator>::value) {
99+
mlvecd.clear();
100+
for (float val : c.mlProbD0()) {
101+
mlvecd.push_back(val);
102+
}
103+
mlvecdbar.clear();
104+
for (float val : c.mlProbD0bar()) {
105+
mlvecdbar.push_back(val);
106+
}
107+
output2ProngTrackmls(cfcollisions.begin().globalIndex(), mlvecd, mlvecdbar);
94108
}
95-
for (float val : c.mlProbD0bar()) {
96-
mlvecdbar.push_back(val);
97-
}
98-
output2ProngTrackmls(cfcollisions.begin().globalIndex(), mlvecd, mlvecdbar);
99109
}
100110
}
101111
}
112+
113+
void processDataML(aod::Collisions::iterator const& cols, aod::BCsWithTimestamps const& bcs, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidatesML const& candidates)
114+
{
115+
processDataT(cols, bcs, cfcollisions, cftracks, candidates);
116+
}
102117
PROCESS_SWITCH(Filter2Prong, processDataML, "Process data D0 candidates with ML", false);
103118

104-
void processData(aod::Collisions::iterator const&, aod::BCsWithTimestamps const&, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidates const& candidates)
119+
void processData(aod::Collisions::iterator const& cols, aod::BCsWithTimestamps const& bcs, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidates const& candidates)
105120
{
106-
if (cfcollisions.size() <= 0 || cftracks.size() <= 0)
107-
return; // rejected collision
108-
if (cfgVerbosity > 0 && candidates.size() > 0)
109-
LOGF(info, "Candidates for collision: %lu, cfcollisions: %lu, CFTracks: %lu", candidates.size(), cfcollisions.size(), cftracks.size());
110-
for (auto& c : candidates) {
111-
int prongCFId[2] = {-1, -1};
112-
for (auto& cftrack : cftracks) {
113-
if (c.prong0Id() == cftrack.trackId()) {
114-
prongCFId[0] = cftrack.globalIndex();
115-
break;
116-
}
117-
}
118-
for (auto& cftrack : cftracks) {
119-
if (c.prong1Id() == cftrack.trackId()) {
120-
prongCFId[1] = cftrack.globalIndex();
121-
break;
122-
}
123-
}
124-
// look-up the collision id
125-
if ((c.hfflag() & (1 << aod::hf_cand_2prong::DecayType::D0ToPiK)) == 0)
126-
continue;
127-
if (cfgYMax >= 0.0f && std::abs(hfHelper.yD0(c)) > cfgYMax)
128-
continue;
129-
if (c.isSelD0() > 0)
130-
output2ProngTracks(cfcollisions.begin().globalIndex(),
131-
prongCFId[0], prongCFId[1], c.pt(), c.eta(), c.phi(), hfHelper.invMassD0ToPiK(c), aod::cf2prongtrack::D0ToPiK);
132-
if (c.isSelD0bar() > 0)
133-
output2ProngTracks(cfcollisions.begin().globalIndex(),
134-
prongCFId[0], prongCFId[1], c.pt(), c.eta(), c.phi(), hfHelper.invMassD0barToKPi(c), aod::cf2prongtrack::D0barToKPi);
135-
}
121+
processDataT(cols, bcs, cfcollisions, cftracks, candidates);
136122
}
137123
PROCESS_SWITCH(Filter2Prong, processData, "Process data D0 candidates", true);
138124
}; // struct

‎PWGCF/Tasks/correlations.cxx‎

Lines changed: 17 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -675,7 +675,8 @@ struct CorrelationTask {
675675
}
676676
PROCESS_SWITCH(CorrelationTask, processSame2ProngDerived, "Process same event on derived data", false);
677677

678-
void processSame2Prong2Prong(DerivedCollisions::iterator const& collision, soa::Filtered<aod::CF2ProngTracks> const& p2tracks)
678+
template <class p2type>
679+
void processSame2Prong2ProngT(DerivedCollisions::iterator const& collision, p2type const& p2tracks)
679680
{
680681
BinningTypeDerived configurableBinningDerived{{axisVertex, axisMultiplicity}, true}; // true is for 'ignore overflows' (true by default). Underflows and overflows will have bin -1.
681682
if (cfgVerbosity > 0) {
@@ -697,29 +698,16 @@ struct CorrelationTask {
697698
fillCorrelations<CorrelationContainer::kCFStepCorrected>(same, p2tracks, p2tracks, multiplicity, collision.posZ(), 0, 1.0f);
698699
}
699700
}
701+
702+
void processSame2Prong2Prong(DerivedCollisions::iterator const& collision, soa::Filtered<aod::CF2ProngTracks> const& p2tracks)
703+
{
704+
processSame2Prong2ProngT(collision, p2tracks);
705+
}
700706
PROCESS_SWITCH(CorrelationTask, processSame2Prong2Prong, "Process same event on derived data", false);
701707

702708
void processSame2Prong2ProngML(DerivedCollisions::iterator const& collision, soa::Filtered<soa::Join<aod::CF2ProngTracks, aod::CF2ProngTrackmls>> const& p2tracks)
703709
{
704-
BinningTypeDerived configurableBinningDerived{{axisVertex, axisMultiplicity}, true}; // true is for 'ignore overflows' (true by default). Underflows and overflows will have bin -1.
705-
if (cfgVerbosity > 0) {
706-
LOGF(info, "processSame2ProngDerived: 2-prong candidates: %d | Vertex: %.1f | Multiplicity/Centrality: %.1f", p2tracks.size(), collision.posZ(), collision.multiplicity());
707-
}
708-
loadEfficiency(collision.timestamp());
709-
710-
const auto multiplicity = collision.multiplicity();
711-
712-
int bin = configurableBinningDerived.getBin({collision.posZ(), collision.multiplicity()});
713-
registry.fill(HIST("eventcount_same"), bin);
714-
fillQA(collision, multiplicity, p2tracks, p2tracks);
715-
716-
same->fillEvent(multiplicity, CorrelationContainer::kCFStepReconstructed);
717-
fillCorrelations<CorrelationContainer::kCFStepReconstructed>(same, p2tracks, p2tracks, multiplicity, collision.posZ(), 0, 1.0f);
718-
719-
if (cfg.mEfficiencyAssociated || cfg.mEfficiencyTrigger) {
720-
same->fillEvent(multiplicity, CorrelationContainer::kCFStepCorrected);
721-
fillCorrelations<CorrelationContainer::kCFStepCorrected>(same, p2tracks, p2tracks, multiplicity, collision.posZ(), 0, 1.0f);
722-
}
710+
processSame2Prong2ProngT(collision, p2tracks);
723711
}
724712
PROCESS_SWITCH(CorrelationTask, processSame2Prong2ProngML, "Process same event on derived data", false);
725713

@@ -847,12 +835,13 @@ struct CorrelationTask {
847835
}
848836
PROCESS_SWITCH(CorrelationTask, processMixed2ProngDerived, "Process mixed events on derived data", false);
849837

850-
void processMixed2Prong2Prong(DerivedCollisions const& collisions, soa::Filtered<aod::CF2ProngTracks> const& p2tracks)
838+
template <class p2type>
839+
void processMixed2Prong2ProngT(DerivedCollisions const& collisions, p2type const& p2tracks)
851840
{
852841
BinningTypeDerived configurableBinningDerived{{axisVertex, axisMultiplicity}, true}; // true is for 'ignore overflows' (true by default). Underflows and overflows will have bin -1.
853842
// Strictly upper categorised collisions, for cfgNoMixedEvents combinations per bin, skipping those in entry -1
854843
auto tracksTuple = std::make_tuple(p2tracks);
855-
SameKindPair<DerivedCollisions, soa::Filtered<aod::CF2ProngTracks>, BinningTypeDerived> pairs{configurableBinningDerived, cfgNoMixedEvents, -1, collisions, tracksTuple, &cache}; // -1 is the number of the bin to skip
844+
SameKindPair<DerivedCollisions, p2type, BinningTypeDerived> pairs{configurableBinningDerived, cfgNoMixedEvents, -1, collisions, tracksTuple, &cache}; // -1 is the number of the bin to skip
856845

857846
for (auto it = pairs.begin(); it != pairs.end(); it++) {
858847
auto& [collision1, tracks1, collision2, tracks2] = *it;
@@ -884,44 +873,16 @@ struct CorrelationTask {
884873
}
885874
}
886875
}
876+
877+
void processMixed2Prong2Prong(DerivedCollisions const& collisions, soa::Filtered<aod::CF2ProngTracks> const& p2tracks)
878+
{
879+
processMixed2Prong2ProngT(collisions, p2tracks);
880+
}
887881
PROCESS_SWITCH(CorrelationTask, processMixed2Prong2Prong, "Process mixed events on derived data", false);
888882

889883
void processMixed2Prong2ProngML(DerivedCollisions const& collisions, soa::Filtered<soa::Join<aod::CF2ProngTracks, aod::CF2ProngTrackmls>> const& p2tracks)
890884
{
891-
BinningTypeDerived configurableBinningDerived{{axisVertex, axisMultiplicity}, true}; // true is for 'ignore overflows' (true by default). Underflows and overflows will have bin -1.
892-
// Strictly upper categorised collisions, for cfgNoMixedEvents combinations per bin, skipping those in entry -1
893-
auto tracksTuple = std::make_tuple(p2tracks);
894-
SameKindPair<DerivedCollisions, soa::Filtered<soa::Join<aod::CF2ProngTracks, aod::CF2ProngTrackmls>>, BinningTypeDerived> pairs{configurableBinningDerived, cfgNoMixedEvents, -1, collisions, tracksTuple, &cache}; // -1 is the number of the bin to skip
895-
896-
for (auto it = pairs.begin(); it != pairs.end(); it++) {
897-
auto& [collision1, tracks1, collision2, tracks2] = *it;
898-
int bin = configurableBinningDerived.getBin({collision1.posZ(), collision1.multiplicity()});
899-
float eventWeight = 1.0f / it.currentWindowNeighbours();
900-
int field = 0;
901-
if (cfgTwoTrackCut > 0) {
902-
field = getMagneticField(collision1.timestamp());
903-
}
904-
905-
if (cfgVerbosity > 0) {
906-
LOGF(info, "processMixedDerived: Mixed collisions bin: %d pair: [%d, %d] %d (%.3f, %.3f), %d (%.3f, %.3f)", bin, it.isNewWindow(), it.currentWindowNeighbours(), collision1.globalIndex(), collision1.posZ(), collision1.multiplicity(), collision2.globalIndex(), collision2.posZ(), collision2.multiplicity());
907-
}
908-
909-
if (it.isNewWindow()) {
910-
loadEfficiency(collision1.timestamp());
911-
mixed->fillEvent(collision1.multiplicity(), CorrelationContainer::kCFStepReconstructed);
912-
}
913-
914-
// LOGF(info, "Tracks: %d and %d entries", tracks1.size(), tracks2.size());
915-
916-
registry.fill(HIST("eventcount_mixed"), bin);
917-
fillCorrelations<CorrelationContainer::kCFStepReconstructed>(mixed, tracks1, tracks2, collision1.multiplicity(), collision1.posZ(), field, eventWeight);
918-
if (cfg.mEfficiencyAssociated || cfg.mEfficiencyTrigger) {
919-
if (it.isNewWindow()) {
920-
mixed->fillEvent(collision1.multiplicity(), CorrelationContainer::kCFStepCorrected);
921-
}
922-
fillCorrelations<CorrelationContainer::kCFStepCorrected>(mixed, tracks1, tracks2, collision1.multiplicity(), collision1.posZ(), field, eventWeight);
923-
}
924-
}
885+
processMixed2Prong2ProngT(collisions, p2tracks);
925886
}
926887
PROCESS_SWITCH(CorrelationTask, processMixed2Prong2ProngML, "Process mixed events on derived data", false);
927888

0 commit comments

Comments
 (0)