Skip to content

Commit c70d079

Browse files
mario-kruegermario-krueger
committed
improve PCC performance
1 parent dbac26b commit c70d079

File tree

1 file changed

+106
-97
lines changed

1 file changed

+106
-97
lines changed

‎PWGLF/TableProducer/Nuspex/particleCompositionCorrection.cxx‎

Lines changed: 106 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -49,15 +49,14 @@ struct ParticleCompositionCorrection {
4949
- support collision systems beyond pp
5050
- add QA task illustrating improved mc/data matching of DCA distributions after scaling of secondaries
5151
- 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
5452
*/
5553

5654
Service<o2::framework::O2DatabasePDG> pdg;
5755
o2::ccdb::CcdbApi ccdbApi;
5856

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

76+
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);
77+
7778
HistogramRegistry histos;
7879
};
7980

@@ -108,7 +109,7 @@ void ParticleCompositionCorrection::init(InitContext const&)
108109
histos.add("frac/data/kaon", "", kTProfile, {ptAxis});
109110
histos.add("frac/data/proton", "", kTProfile, {ptAxis});
110111
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...
112+
histos.addClone("frac/data/", "frac/mc/");
112113

113114
histos.add("weight/pion", "", kTProfile, {ptAxis});
114115
histos.add("weight/kaon", "", kTProfile, {ptAxis});
@@ -120,6 +121,105 @@ void ParticleCompositionCorrection::init(InitContext const&)
120121
}
121122
}
122123

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

141-
// fill particle composition table with corresponding weights
241+
std::map<int32_t, std::tuple<float, float, float>> storedWeights;
142242
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
243+
auto [weight, weightSysUp, weightSysDown] = getWeights(particles, particle, storedWeights, dNdEta);
244+
pccTable(weight, weightSysUp, weightSysDown);
236245
}
237246
}

0 commit comments

Comments
 (0)