Skip to content

Commit 6299095

Browse files
mario-kruegermario-krueger
committed
add particle-composition correction
1 parent 894709e commit 6299095

File tree

3 files changed

+261
-0
lines changed

3 files changed

+261
-0
lines changed

‎PWGLF/DataModel/particleCompositionCorrectionTable.h‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
2+
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3+
// All rights not expressly granted are reserved.
4+
//
5+
// This software is distributed under the terms of the GNU General Public
6+
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7+
//
8+
// In applying this license CERN does not waive the privileges and immunities
9+
// granted to it by virtue of its status as an Intergovernmental Organization
10+
// or submit itself to any jurisdiction.
11+
12+
///
13+
/// \file ParticleCompositionCorrectionTable.h
14+
/// \brief Table for scaling MC particle abundances to match measurements
15+
/// \author Mario Krüger <mario.kruger@cern.ch>
16+
///
17+
18+
#ifndef PWGLF_DATAMODEL_PARTICLECOMPOSITIONCORRECTION_H_
19+
#define PWGLF_DATAMODEL_PARTICLECOMPOSITIONCORRECTION_H_
20+
21+
#include "Framework/ASoA.h"
22+
#include "Framework/AnalysisDataModel.h"
23+
24+
namespace o2::aod
25+
{
26+
namespace PCC
27+
{
28+
DECLARE_SOA_COLUMN(Pccweight, pccweight, float);
29+
DECLARE_SOA_COLUMN(Pccweightsysup, pccweightsysup, float);
30+
DECLARE_SOA_COLUMN(Pccweightsysdown, pccweightsysdown, float);
31+
} // namespace PCC
32+
DECLARE_SOA_TABLE(ParticleCompositionCorrection, "AOD", "PARTICLECOMPOSITIONCORRECTION", PCC::Pccweight, PCC::Pccweightsysup, PCC::Pccweightsysdown);
33+
} // namespace o2::aod
34+
35+
#endif // PWGLF_DATAMODEL_PARTICLECOMPOSITIONCORRECTION_H_

‎PWGLF/TableProducer/Nuspex/CMakeLists.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,3 +108,8 @@ o2physics_add_dpl_workflow(nuclei-antineutron-cex
108108
SOURCES nucleiAntineutronCex.cxx
109109
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
110110
COMPONENT_NAME Analysis)
111+
112+
o2physics_add_dpl_workflow(particle-composition-correction
113+
SOURCES particleCompositionCorrection.cxx
114+
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore
115+
COMPONENT_NAME Analysis)

‎PWGLF/TableProducer/Nuspex/particleCompositionCorrection.cxx‎

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
1+
// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
2+
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3+
// All rights not expressly granted are reserved.
4+
//
5+
// This software is distributed under the terms of the GNU General Public
6+
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7+
//
8+
// In applying this license CERN does not waive the privileges and immunities
9+
// granted to it by virtue of its status as an Intergovernmental Organization
10+
// or submit itself to any jurisdiction.
11+
12+
/// \file particleCompositionCorrection.cxx
13+
/// \brief Task to generate a table of dNdEta, pt and PID dependent weigths for MC particles to reflect the measured particle abundances
14+
/// \author Mario Krüger <mario.kruger@cern.ch>
15+
16+
#include <PWGLF/DataModel/particleCompositionCorrectionTable.h>
17+
18+
#include "Tools/ML/model.h"
19+
20+
#include <Framework/AnalysisTask.h>
21+
#include <Framework/HistogramRegistry.h>
22+
#include <Framework/O2DatabasePDGPlugin.h>
23+
#include <Framework/runDataProcessing.h>
24+
25+
#include <TMCProcess.h>
26+
#include <TPDGCode.h>
27+
28+
#include <map>
29+
#include <string>
30+
#include <vector>
31+
32+
using namespace o2;
33+
using namespace o2::framework;
34+
using namespace o2::ml;
35+
36+
struct ParticleCompositionCorrection {
37+
// produces table of (dNdEta, pt, PID)-dependent weights for adjusting the particle composition in MC to better match the measured particle chemistry
38+
// applying these weights in the analysis allows mitigating on event-by-event basis the corresponding biases in efficiency and contamination of unidentified particles
39+
// weights are determined using measured pi,K,p,labda(as basis for sigma) spectra togehter with their counterpart from pythia simulations
40+
// they are interpolated in dNdEta and pt dimension using DNNs, which then provide particle fractions in data and MC and are stored as .onnx in the CCDB
41+
// weigths are assigned to the primary generated particle as well as its daughter particles (both from decay and material interactions)
42+
// weights are calculated only for particles within the configured kineamatic range and only for a distinct set of mother particles (see code)
43+
// assumes neutral particles require the same scaling as their charged counterparts (e.g. pi0 is scaled the same as pi+)
44+
// multi-strange baryons are assigned scaling factors of the sigma, which should be better than no scaling at all
45+
46+
/*
47+
backlog:
48+
- store final neural networks in ccdb and implement accessing them
49+
- support collision systems beyond pp
50+
- add QA task illustrating improved mc/data matching of DCA distributions after scaling of secondaries
51+
- 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
54+
*/
55+
56+
Service<o2::framework::O2DatabasePDG> pdg;
57+
Configurable<float> etaCut{"etaCut", 0.8f, "eta cut"};
58+
Configurable<float> ptMinCut{"ptMinCut", 0.15f, "pt min cut"};
59+
Configurable<float> ptMaxCut{"ptMaxCut", 10.f, "pt max cut"};
60+
Configurable<bool> skipSec{"skipSec", false, "dont calculate weights for secondaries"};
61+
Configurable<bool> enableQAHistos{"enableQAHistos", true, "enable qa histograms showing the effect of the PCC"};
62+
63+
Configurable<std::string> modelNameData{"modelNameData", "ParticleFractions_pp_data.onnx", "Path to the .onnx file containing the particle fractions in data"};
64+
Configurable<std::string> modelNameMC{"modelNameMC", "ParticleFractions_pp_pythia.onnx", "Path to the .onnx file containing the particle fractions in MC"};
65+
66+
OnnxModel particleFractionsData;
67+
OnnxModel particleFractionsMC;
68+
69+
Produces<o2::aod::ParticleCompositionCorrection> pccTable;
70+
void init(InitContext const& cfgc);
71+
void process(aod::McCollisions::iterator const& mcCollision, aod::McParticles const& particles);
72+
73+
HistogramRegistry histos;
74+
};
75+
76+
WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
77+
{
78+
return WorkflowSpec{adaptAnalysisTask<ParticleCompositionCorrection>(cfgc)};
79+
}
80+
81+
void ParticleCompositionCorrection::init(InitContext const&)
82+
{
83+
particleFractionsData.initModel(modelNameData.value, true);
84+
particleFractionsMC.initModel(modelNameMC.value, true);
85+
86+
if (enableQAHistos) {
87+
std::vector<double> ptBinEdges = {0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75,
88+
0.8, 0.85, 0.9, 0.95, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9,
89+
2.0, 2.2, 2.4, 2.6, 2.8, 3.0, 3.2, 3.4, 3.6, 3.8, 4.0, 4.5, 5.0, 5.5,
90+
6.0, 6.5, 7.0, 8.0, 9.0, 10.0};
91+
const AxisSpec ptAxis{ptBinEdges, "#it{p}_{T} (GeV/#it{c})", "pt"};
92+
93+
histos.add("frac/data/pion", "", kTProfile, {ptAxis});
94+
histos.add("frac/data/kaon", "", kTProfile, {ptAxis});
95+
histos.add("frac/data/proton", "", kTProfile, {ptAxis});
96+
histos.add("frac/data/sigma", "", kTProfile, {ptAxis});
97+
histos.addClone("frac/data/", "frac/mc/"); // FIXME: bug in writer/registry moves one histogram out of subfolder...
98+
99+
histos.add("weight/pion", "", kTProfile, {ptAxis});
100+
histos.add("weight/kaon", "", kTProfile, {ptAxis});
101+
histos.add("weight/proton", "", kTProfile, {ptAxis});
102+
histos.add("weight/sigma", "", kTProfile, {ptAxis});
103+
104+
histos.add("weight/secDec", "", kTProfile, {ptAxis});
105+
histos.add("weight/secMat", "", kTProfile, {ptAxis});
106+
}
107+
}
108+
109+
void ParticleCompositionCorrection::process(aod::McCollisions::iterator const&, aod::McParticles const& particles)
110+
{
111+
// determine dNdEta of the collision
112+
float dNdEta = 0.f;
113+
for (const auto& particle : particles) {
114+
if (!particle.isPhysicalPrimary()) {
115+
continue;
116+
}
117+
auto pdgParticle = pdg->GetParticle(particle.pdgCode());
118+
if (!pdgParticle || pdgParticle->Charge() == 0.) {
119+
continue;
120+
}
121+
if (std::abs(particle.eta()) >= 0.5) { // o2-linter: disable=magic-number (particle density at mid-rapidity)
122+
continue;
123+
}
124+
++dNdEta;
125+
}
126+
127+
// fill particle composition table with corresponding weights
128+
for (const auto& particle : particles) {
129+
float weight = 1.f;
130+
131+
// calculate weights only inside the configured kinematic range
132+
if (std::abs(particle.eta()) < etaCut && particle.pt() > ptMinCut && particle.pt() < ptMaxCut) {
133+
auto refParticleID = particle.index();
134+
135+
// find initial particle of secondaries from decays or interactions with material
136+
while (!particles.iteratorAt(refParticleID).producedByGenerator() && particles.iteratorAt(refParticleID).has_mothers()) {
137+
if (skipSec) {
138+
break;
139+
}
140+
const auto& mother = particle.mothers_first_as<aod::McParticles>();
141+
auto motherID = mother.globalIndex() - particles.offset();
142+
if (motherID == static_cast<uint64_t>(refParticleID)) {
143+
// FIXME: this needs to be investigated!
144+
// LOGP(error, " - !!! particle is its own mother: {} (pdg: {}) status: {} ; process: {}", motherID, mother.pdgCode(), mother.getGenStatusCode(), mother.getProcess());
145+
break;
146+
}
147+
refParticleID = motherID;
148+
}
149+
150+
if (const auto& refParticle = particles.iteratorAt(refParticleID); refParticle.producedByGenerator()) {
151+
152+
auto absPDGCode = std::abs(refParticle.pdgCode());
153+
// translate abs PDG code to PID variable of neural networks (0: pion, 1: kaon, 2: proton, 3: sigma)
154+
static const std::map<int, float> mapPID = {
155+
{PDG_t::kPiPlus, 0.f},
156+
{PDG_t::kPi0, 0.f},
157+
{PDG_t::kKPlus, 1.f},
158+
{PDG_t::kK0Short, 1.f},
159+
{PDG_t::kK0Long, 1.f},
160+
{PDG_t::kProton, 2.f},
161+
{PDG_t::kNeutron, 2.f},
162+
{PDG_t::kSigmaPlus, 3.f},
163+
{PDG_t::kSigmaMinus, 3.f},
164+
{PDG_t::kLambda0, 3.f},
165+
{PDG_t::kSigma0, 3.f},
166+
{PDG_t::kXiMinus, 3.f},
167+
// TODO: potentially extend by xi0/eta/omega/rho/phi/Delta...
168+
};
169+
170+
if (auto iterMapPID = mapPID.find(absPDGCode); iterMapPID != mapPID.end()) {
171+
// LOGP(info, "scaling a {} with status code {} from process {}", particle.pdgCode(), particle.getGenStatusCode(), particle.getProcess());
172+
float pt = refParticle.pt();
173+
174+
// calculate particle fractions and corresponding weight for given reference particle
175+
std::vector<std::vector<float>> input = {{dNdEta}, {pt}, {iterMapPID->second}};
176+
float fracData = particleFractionsData.evalModel(input)[0];
177+
float fracMC = particleFractionsMC.evalModel(input)[0];
178+
if (fracMC) {
179+
weight = fracData / fracMC;
180+
}
181+
182+
if (enableQAHistos && std::abs(particle.eta()) < 0.8) { // o2-linter: disable=magic-number (usual range of charged-partilce measurements)
183+
if (particle.isPhysicalPrimary()) {
184+
if (iterMapPID->first == PDG_t::kPiPlus) {
185+
histos.fill(HIST("frac/data/pion"), pt, fracData);
186+
histos.fill(HIST("frac/mc/pion"), pt, fracMC);
187+
histos.fill(HIST("weight/pion"), pt, weight);
188+
}
189+
if (iterMapPID->first == PDG_t::kKPlus) {
190+
histos.fill(HIST("frac/data/kaon"), pt, fracData);
191+
histos.fill(HIST("frac/mc/kaon"), pt, fracMC);
192+
histos.fill(HIST("weight/kaon"), pt, weight);
193+
}
194+
if (iterMapPID->first == PDG_t::kProton) {
195+
histos.fill(HIST("frac/data/proton"), pt, fracData);
196+
histos.fill(HIST("frac/mc/proton"), pt, fracMC);
197+
histos.fill(HIST("weight/proton"), pt, weight);
198+
}
199+
if (iterMapPID->first == PDG_t::kSigmaPlus || iterMapPID->first == PDG_t::kSigmaMinus) {
200+
histos.fill(HIST("frac/data/sigma"), pt, fracData);
201+
histos.fill(HIST("frac/mc/sigma"), pt, fracMC);
202+
histos.fill(HIST("weight/sigma"), pt, weight);
203+
}
204+
} else {
205+
auto pdgParticle = pdg->GetParticle(particle.pdgCode());
206+
if (pdgParticle && pdgParticle->Charge() != 0.) {
207+
// fill the decay daughter info only (usually there is max. one parent generation so no relevant double counting here)
208+
if (particle.getProcess() == TMCProcess::kPDecay) {
209+
histos.fill(HIST("weight/secDec"), pt, weight);
210+
} else if (particle.getProcess() == TMCProcess::kPHInhelastic) {
211+
histos.fill(HIST("weight/secMat"), pt, weight);
212+
}
213+
}
214+
}
215+
}
216+
}
217+
}
218+
}
219+
pccTable(weight, weight, weight); // TODO: put here systematic variations of the weights
220+
}
221+
}

0 commit comments

Comments
 (0)