Skip to content

Commit bd9fae8

Browse files
mario-kruegermario-krueger
authored
[PWGLF] improve PCC performance (#13468)
1 parent a2f03fd commit bd9fae8

File tree

1 file changed

+107
-97
lines changed

1 file changed

+107
-97
lines changed

‎PWGLF/TableProducer/Nuspex/particleCompositionCorrection.cxx‎

Lines changed: 107 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828

2929
#include <map>
3030
#include <string>
31+
#include <tuple>
3132
#include <vector>
3233

3334
using namespace o2;
@@ -49,15 +50,14 @@ struct ParticleCompositionCorrection {
4950
- support collision systems beyond pp
5051
- add QA task illustrating improved mc/data matching of DCA distributions after scaling of secondaries
5152
- extend PCC weight table by columns with systematic variations (up/down)
52-
- investigate why particles can be their own mothers
53-
- investigate why histogram registry writing does not respect subfolder structure
5453
*/
5554

5655
Service<o2::framework::O2DatabasePDG> pdg;
5756
o2::ccdb::CcdbApi ccdbApi;
5857

5958
Configurable<bool> skipAll{"skipAll", false, "run table producer in dummy mode, i.e. skip all computations and fill with 1"};
6059
Configurable<bool> skipSec{"skipSec", false, "dont calculate weights for secondaries"};
60+
Configurable<bool> skipNonPhysicalPrim{"skipNonPhysicalPrim", true, "dont calculate weights for particles that are not (originating from) physical primaries; i.e. reject (decays of) pi0 etc."};
6161
Configurable<float> etaCut{"etaCut", 0.8f, "eta cut"};
6262
Configurable<float> ptMinCut{"ptMinCut", 0.15f, "pt min cut"};
6363
Configurable<float> ptMaxCut{"ptMaxCut", 10.f, "pt max cut"};
@@ -74,6 +74,8 @@ struct ParticleCompositionCorrection {
7474
void init(InitContext const& cfgc);
7575
void process(aod::McCollisions::iterator const& mcCollision, aod::McParticles const& particles);
7676

77+
std::tuple<float, float, float> getWeights(aod::McParticles const& particles, aod::McParticles::iterator const& particle, std::map<int32_t, std::tuple<float, float, float>>& storedWeights, float dNdEta);
78+
7779
HistogramRegistry histos;
7880
};
7981

@@ -108,7 +110,7 @@ void ParticleCompositionCorrection::init(InitContext const&)
108110
histos.add("frac/data/kaon", "", kTProfile, {ptAxis});
109111
histos.add("frac/data/proton", "", kTProfile, {ptAxis});
110112
histos.add("frac/data/sigma", "", kTProfile, {ptAxis});
111-
histos.addClone("frac/data/", "frac/mc/"); // FIXME: bug in writer/registry moves one histogram out of subfolder...
113+
histos.addClone("frac/data/", "frac/mc/");
112114

113115
histos.add("weight/pion", "", kTProfile, {ptAxis});
114116
histos.add("weight/kaon", "", kTProfile, {ptAxis});
@@ -120,6 +122,105 @@ void ParticleCompositionCorrection::init(InitContext const&)
120122
}
121123
}
122124

125+
std::tuple<float, float, float> ParticleCompositionCorrection::getWeights(aod::McParticles const& particles, aod::McParticles::iterator const& particle, std::map<int32_t, std::tuple<float, float, float>>& storedWeights, float dNdEta)
126+
{
127+
static const std::tuple<float, float, float> noWeights = {1.f, 1.f, 1.f};
128+
129+
if (skipAll || std::abs(particle.eta()) > etaCut || particle.pt() < ptMinCut || particle.pt() > ptMaxCut) {
130+
return noWeights;
131+
}
132+
133+
if (particle.producedByGenerator()) {
134+
if (skipNonPhysicalPrim && !particle.isPhysicalPrimary()) {
135+
return noWeights;
136+
}
137+
auto absPDGCode = std::abs(particle.pdgCode());
138+
// translate abs PDG code to PID variable of neural networks (0: pion, 1: kaon, 2: proton, 3: sigma)
139+
static const std::map<int, float> mapPID = {
140+
{PDG_t::kPiPlus, 0.f},
141+
{PDG_t::kPi0, 0.f},
142+
{PDG_t::kKPlus, 1.f},
143+
{PDG_t::kK0Short, 1.f},
144+
{PDG_t::kK0Long, 1.f},
145+
{PDG_t::kProton, 2.f},
146+
{PDG_t::kNeutron, 2.f},
147+
{PDG_t::kSigmaPlus, 3.f},
148+
{PDG_t::kSigmaMinus, 3.f},
149+
{PDG_t::kLambda0, 3.f},
150+
{PDG_t::kSigma0, 3.f},
151+
{PDG_t::kXiMinus, 3.f},
152+
// TODO: potentially extend by xi0/eta/omega/rho/phi/Delta...
153+
// pdg codes defined in AliceO2/Common/Constants/include/CommonConstants/PhysicsConstants.h
154+
// e.g. o2::constants::physics::Pdg::kEta
155+
};
156+
157+
if (auto iterMapPID = mapPID.find(absPDGCode); iterMapPID != mapPID.end()) {
158+
// LOGP(info, "scaling a {} with status code {} from process {}", particle.pdgCode(), particle.getGenStatusCode(), particle.getProcess());
159+
float pt = particle.pt();
160+
161+
// calculate particle fractions and corresponding weight for given reference particle
162+
std::vector<std::vector<float>> input = {{dNdEta}, {pt}, {iterMapPID->second}};
163+
float fracData = particleFractionsData.evalModel(input)[0];
164+
float fracMC = particleFractionsMC.evalModel(input)[0];
165+
float weight = (fracMC) ? fracData / fracMC : 1.f;
166+
std::tuple<float, float, float> weights = {weight, weight, weight};
167+
if (!skipSec && particle.has_daughters()) {
168+
storedWeights[particle.index()] = weights;
169+
}
170+
if (enableQAHistos && particle.isPhysicalPrimary() && std::abs(particle.eta()) < 0.8) { // o2-linter: disable=magic-number (usual range of charged-partilce measurements)
171+
if (iterMapPID->first == PDG_t::kPiPlus) {
172+
histos.fill(HIST("frac/data/pion"), pt, fracData);
173+
histos.fill(HIST("frac/mc/pion"), pt, fracMC);
174+
histos.fill(HIST("weight/pion"), pt, weight);
175+
}
176+
if (iterMapPID->first == PDG_t::kKPlus) {
177+
histos.fill(HIST("frac/data/kaon"), pt, fracData);
178+
histos.fill(HIST("frac/mc/kaon"), pt, fracMC);
179+
histos.fill(HIST("weight/kaon"), pt, weight);
180+
}
181+
if (iterMapPID->first == PDG_t::kProton) {
182+
histos.fill(HIST("frac/data/proton"), pt, fracData);
183+
histos.fill(HIST("frac/mc/proton"), pt, fracMC);
184+
histos.fill(HIST("weight/proton"), pt, weight);
185+
}
186+
if (iterMapPID->first == PDG_t::kSigmaPlus || iterMapPID->first == PDG_t::kSigmaMinus) {
187+
histos.fill(HIST("frac/data/sigma"), pt, fracData);
188+
histos.fill(HIST("frac/mc/sigma"), pt, fracMC);
189+
histos.fill(HIST("weight/sigma"), pt, weight);
190+
}
191+
}
192+
return weights;
193+
}
194+
} else if (!skipSec) {
195+
auto refParticleID = particle.index();
196+
// LOGP(error, "Particle [{}] {} from process {}", refParticleID, particle.pdgCode(), particle.getProcess());
197+
while (!particles.iteratorAt(refParticleID).producedByGenerator() && particles.iteratorAt(refParticleID).has_mothers()) {
198+
auto motherID = particles.iteratorAt(refParticleID).mothersIds()[0] - particles.offset();
199+
// LOGP(error, "-> mom [{}] {} from process {}", motherID, particles.iteratorAt(motherID).pdgCode(), particles.iteratorAt(motherID).getProcess());
200+
refParticleID = motherID;
201+
}
202+
203+
if (storedWeights.find(refParticleID) == storedWeights.end()) {
204+
// LOGP(error, " no ref particle stored for particle {} from process {}!!", particle.pdgCode(), particle.getProcess());
205+
return noWeights;
206+
}
207+
208+
if (enableQAHistos) {
209+
float weight = get<0>(storedWeights.at(refParticleID));
210+
auto pdgParticle = pdg->GetParticle(particle.pdgCode());
211+
if (pdgParticle && pdgParticle->Charge() != 0.) {
212+
if (particle.getProcess() == TMCProcess::kPDecay) {
213+
histos.fill(HIST("weight/secDec"), particle.pt(), weight);
214+
} else if (particle.getProcess() == TMCProcess::kPHInhelastic || particle.getProcess() == TMCProcess::kPHadronic || particle.getProcess() == TMCProcess::kPHElastic) {
215+
histos.fill(HIST("weight/secMat"), particle.pt(), weight);
216+
}
217+
}
218+
}
219+
return storedWeights.at(refParticleID);
220+
}
221+
return noWeights;
222+
}
223+
123224
void ParticleCompositionCorrection::process(aod::McCollisions::iterator const&, aod::McParticles const& particles)
124225
{
125226
// determine dNdEta of the collision
@@ -138,100 +239,9 @@ void ParticleCompositionCorrection::process(aod::McCollisions::iterator const&,
138239
++dNdEta;
139240
}
140241

141-
// fill particle composition table with corresponding weights
242+
std::map<int32_t, std::tuple<float, float, float>> storedWeights;
142243
for (const auto& particle : particles) {
143-
float weight = 1.f;
144-
145-
// calculate weights only inside the configured kinematic range
146-
if (!skipAll && std::abs(particle.eta()) < etaCut && particle.pt() > ptMinCut && particle.pt() < ptMaxCut) {
147-
auto refParticleID = particle.index();
148-
149-
// find initial particle of secondaries from decays or interactions with material
150-
while (!particles.iteratorAt(refParticleID).producedByGenerator() && particles.iteratorAt(refParticleID).has_mothers()) {
151-
if (skipSec) {
152-
break;
153-
}
154-
const auto& mother = particle.mothers_first_as<aod::McParticles>();
155-
auto motherID = mother.globalIndex() - particles.offset();
156-
if (motherID == static_cast<uint64_t>(refParticleID)) {
157-
// FIXME: this needs to be investigated!
158-
// LOGP(error, " - !!! particle is its own mother: {} (pdg: {}) status: {} ; process: {}", motherID, mother.pdgCode(), mother.getGenStatusCode(), mother.getProcess());
159-
break;
160-
}
161-
refParticleID = motherID;
162-
}
163-
164-
if (const auto& refParticle = particles.iteratorAt(refParticleID); refParticle.producedByGenerator()) {
165-
166-
auto absPDGCode = std::abs(refParticle.pdgCode());
167-
// translate abs PDG code to PID variable of neural networks (0: pion, 1: kaon, 2: proton, 3: sigma)
168-
static const std::map<int, float> mapPID = {
169-
{PDG_t::kPiPlus, 0.f},
170-
{PDG_t::kPi0, 0.f},
171-
{PDG_t::kKPlus, 1.f},
172-
{PDG_t::kK0Short, 1.f},
173-
{PDG_t::kK0Long, 1.f},
174-
{PDG_t::kProton, 2.f},
175-
{PDG_t::kNeutron, 2.f},
176-
{PDG_t::kSigmaPlus, 3.f},
177-
{PDG_t::kSigmaMinus, 3.f},
178-
{PDG_t::kLambda0, 3.f},
179-
{PDG_t::kSigma0, 3.f},
180-
{PDG_t::kXiMinus, 3.f},
181-
// TODO: potentially extend by xi0/eta/omega/rho/phi/Delta...
182-
// pdg codes defined in AliceO2/Common/Constants/include/CommonConstants/PhysicsConstants.h
183-
// e.g. o2::constants::physics::Pdg::kEta
184-
};
185-
186-
if (auto iterMapPID = mapPID.find(absPDGCode); iterMapPID != mapPID.end()) {
187-
// LOGP(info, "scaling a {} with status code {} from process {}", particle.pdgCode(), particle.getGenStatusCode(), particle.getProcess());
188-
float pt = refParticle.pt();
189-
190-
// calculate particle fractions and corresponding weight for given reference particle
191-
std::vector<std::vector<float>> input = {{dNdEta}, {pt}, {iterMapPID->second}};
192-
float fracData = particleFractionsData.evalModel(input)[0];
193-
float fracMC = particleFractionsMC.evalModel(input)[0];
194-
if (fracMC) {
195-
weight = fracData / fracMC;
196-
}
197-
198-
if (enableQAHistos && std::abs(particle.eta()) < 0.8) { // o2-linter: disable=magic-number (usual range of charged-partilce measurements)
199-
if (particle.isPhysicalPrimary()) {
200-
if (iterMapPID->first == PDG_t::kPiPlus) {
201-
histos.fill(HIST("frac/data/pion"), pt, fracData);
202-
histos.fill(HIST("frac/mc/pion"), pt, fracMC);
203-
histos.fill(HIST("weight/pion"), pt, weight);
204-
}
205-
if (iterMapPID->first == PDG_t::kKPlus) {
206-
histos.fill(HIST("frac/data/kaon"), pt, fracData);
207-
histos.fill(HIST("frac/mc/kaon"), pt, fracMC);
208-
histos.fill(HIST("weight/kaon"), pt, weight);
209-
}
210-
if (iterMapPID->first == PDG_t::kProton) {
211-
histos.fill(HIST("frac/data/proton"), pt, fracData);
212-
histos.fill(HIST("frac/mc/proton"), pt, fracMC);
213-
histos.fill(HIST("weight/proton"), pt, weight);
214-
}
215-
if (iterMapPID->first == PDG_t::kSigmaPlus || iterMapPID->first == PDG_t::kSigmaMinus) {
216-
histos.fill(HIST("frac/data/sigma"), pt, fracData);
217-
histos.fill(HIST("frac/mc/sigma"), pt, fracMC);
218-
histos.fill(HIST("weight/sigma"), pt, weight);
219-
}
220-
} else {
221-
auto pdgParticle = pdg->GetParticle(particle.pdgCode());
222-
if (pdgParticle && pdgParticle->Charge() != 0.) {
223-
// fill the decay daughter info only (usually there is max. one parent generation so no relevant double counting here)
224-
if (particle.getProcess() == TMCProcess::kPDecay) {
225-
histos.fill(HIST("weight/secDec"), pt, weight);
226-
} else if (particle.getProcess() == TMCProcess::kPHInhelastic) {
227-
histos.fill(HIST("weight/secMat"), pt, weight);
228-
}
229-
}
230-
}
231-
}
232-
}
233-
}
234-
}
235-
pccTable(weight, weight, weight); // TODO: put here systematic variations of the weights
244+
auto [weight, weightSysUp, weightSysDown] = getWeights(particles, particle, storedWeights, dNdEta);
245+
pccTable(weight, weightSysUp, weightSysDown);
236246
}
237247
}

0 commit comments

Comments
 (0)