Skip to content

Commit f522667

Browse files
mario-kruegermario-krueger
committed
add particle composition correction
1 parent 28dfc98 commit f522667

File tree

3 files changed

+250
-0
lines changed

3 files changed

+250
-0
lines changed

‎PWGLF/DataModel/particleCompositionCorrectionTable.h‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
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(Weight, weight, float);
29+
}
30+
DECLARE_SOA_TABLE(ParticleCompositionCorrection, "AOD", "PARTICLECOMPOSITIONCORRECTION", PCC::Weight);
31+
} // namespace o2::aod
32+
33+
#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: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,212 @@
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 "Common/Core/MetadataHelper.h"
17+
#include "Common/Core/TrackSelection.h"
18+
#include "Common/Core/TrackSelectionDefaults.h"
19+
#include "Common/DataModel/Centrality.h"
20+
#include "Common/DataModel/EventSelection.h"
21+
#include "Common/DataModel/TrackSelectionTables.h"
22+
23+
#include <Framework/AnalysisTask.h>
24+
#include <Framework/HistogramRegistry.h>
25+
#include <Framework/O2DatabasePDGPlugin.h>
26+
#include <Framework/runDataProcessing.h>
27+
#include <ReconstructionDataFormats/Track.h>
28+
29+
#include <PWGLF/DataModel/particleCompositionCorrectionTable.h>
30+
#include "Tools/ML/model.h"
31+
32+
#include <unordered_set>
33+
#include <vector>
34+
#include <TPDGCode.h>
35+
#include <TMCProcess.h>
36+
37+
using namespace o2;
38+
using namespace o2::framework;
39+
using namespace o2::ml;
40+
41+
struct ParticleCompositionCorrection {
42+
// produces table of (dNdEta, pt, PID)-dependent weights for adjusting the particle composition in MC to better match the measured particle chemistry
43+
// weights are determined using measured pi,K,p,labda (to construct sigma) spectra togehter with their counterpart from pythia simulations
44+
// 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
45+
// weigths are assigned to the primary generated particle as well as its daughter particles (both from decay and material interactions)
46+
// weights are calculated only for particles within the configured kineamatic range and only for a distinct set of mother particles (see code)
47+
// assumes neutral particles require the same scaling as their charged counterparts (e.g. pi0 is scaled the same as pi+)
48+
// multi-strange baryons are assigned scaling factors of the sigma, which should be better than no scaling at all
49+
// applying the weights in the analysis allows mitigating biases in inclusive charged particle efficiency and contamination originating from the wrong particle composition in the generator on event-by-event basis
50+
51+
/*
52+
backlog:
53+
- store final neural networks in ccdb and implement accessing them
54+
- support collision systems beyond pp
55+
- add QA task illustrating improved mc/data matching of DCA distributions after scaling of secondaries
56+
- extend PCC weight table by columns with systematic variations (up/down)
57+
- investigate why particles can be their own mothers
58+
- investigate why histogram registry writing does not respect subfolder structure
59+
*/
60+
61+
Service<o2::framework::O2DatabasePDG> pdg;
62+
Configurable<float> etaCut{"etaCut", 0.8f, "eta cut"};
63+
Configurable<float> ptMinCut{"ptMinCut", 0.15f, "pt min cut"};
64+
Configurable<float> ptMaxCut{"ptMaxCut", 10.f, "pt max cut"};
65+
Configurable<bool> skipSec{"skipSec", true, "dont calculate weights for secondaries"};
66+
Configurable<bool> enableQAHistos{"enableQAHistos", true, "enable qa histograms showing the effect of the PCC"};
67+
68+
Configurable<std::string> modelNameData{"particleFractionsModelData", "ParticleFractions_pp_data.onnx", "Path to the local .onnx file"};
69+
Configurable<std::string> modelNameMC{"particleFractionsModelMC", "ParticleFractions_pp_pythia.onnx", "Path to the local .onnx file"};
70+
71+
OnnxModel particleFractionsData;
72+
OnnxModel particleFractionsMC;
73+
74+
Produces<o2::aod::ParticleCompositionCorrection> pccTable;
75+
void init(InitContext const& cfgc);
76+
void process(aod::McCollisions::iterator const& mcCollision, aod::McParticles const& particles);
77+
78+
HistogramRegistry histos;
79+
};
80+
81+
WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
82+
{
83+
return WorkflowSpec{adaptAnalysisTask<ParticleCompositionCorrection>(cfgc)};
84+
}
85+
86+
void ParticleCompositionCorrection::init(InitContext const&)
87+
{
88+
particleFractionsData.initModel(modelNameData.value, true);
89+
particleFractionsMC.initModel(modelNameMC.value, true);
90+
91+
if (enableQAHistos) {
92+
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,
93+
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,
94+
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,
95+
6.0, 6.5, 7.0, 8.0, 9.0, 10.0};
96+
const AxisSpec ptAxis{ptBinEdges, "#it{p}_{T} (GeV/#it{c})", "pt"};
97+
98+
histos.add("frac/data/pion", "", kTProfile, {ptAxis});
99+
histos.add("frac/data/kaon", "", kTProfile, {ptAxis});
100+
histos.add("frac/data/proton", "", kTProfile, {ptAxis});
101+
histos.add("frac/data/sigma", "", kTProfile, {ptAxis});
102+
histos.addClone("frac/data/", "frac/mc/"); // FIXME: bug in writer/registry moves one histogram out of subfolder...
103+
104+
histos.add("weight/pion", "", kTProfile, {ptAxis});
105+
histos.add("weight/kaon", "", kTProfile, {ptAxis});
106+
histos.add("weight/proton", "", kTProfile, {ptAxis});
107+
histos.add("weight/sigma", "", kTProfile, {ptAxis});
108+
109+
histos.add("weight/secDec", "", kTProfile, {ptAxis});
110+
histos.add("weight/secMat", "", kTProfile, {ptAxis});
111+
}
112+
}
113+
114+
void ParticleCompositionCorrection::process(aod::McCollisions::iterator const&, aod::McParticles const& particles)
115+
{
116+
// determine dNdEta of the collision
117+
float dNdEta = 0.f;
118+
for (const auto& particle : particles) {
119+
if (!particle.isPhysicalPrimary()) continue;
120+
auto pdgParticle = pdg->GetParticle(particle.pdgCode());
121+
if (!pdgParticle || pdgParticle->Charge() == 0.) continue;
122+
if (std::abs(particle.eta()) >= 0.5) continue;
123+
++dNdEta;
124+
}
125+
126+
// fill particle composition table with corresponding weights
127+
for (const auto& particle : particles) {
128+
float weight = 1.f;
129+
130+
// calculate weights only inside the configured kinematic range
131+
if (std::abs(particle.eta()) < etaCut && particle.pt() > ptMinCut && particle.pt() < ptMaxCut) {
132+
auto refParticleID = particle.index();
133+
134+
// find initial particle of secondaries from decays or interactions with material
135+
while (!particles.iteratorAt(refParticleID).producedByGenerator() && particles.iteratorAt(refParticleID).has_mothers()) {
136+
if (skipSec) break;
137+
const auto& mother = particle.mothers_first_as<aod::McParticles>();
138+
auto motherID = mother.globalIndex() - particles.offset();
139+
if (motherID == static_cast<long unsigned int>(refParticleID)) { // FIXME: this needs to be investigated!
140+
// LOGP(error, " - !!! particle is its own mother: {} (pdg: {}) status: {} ; process: {}", motherID, mother.pdgCode(), mother.getGenStatusCode(), mother.getProcess());
141+
break;
142+
}
143+
refParticleID = motherID;
144+
}
145+
146+
if (const auto& refParticle = particles.iteratorAt(refParticleID); refParticle.producedByGenerator()) {
147+
148+
auto absPDGCode = std::abs(refParticle.pdgCode());
149+
// translate abs PDG code to PID variable of neural networks (0: pion, 1: kaon, 2: proton, 3: sigma)
150+
static const std::map<int, float> mapPID = {
151+
{PDG_t::kPiPlus, 0.f},
152+
{PDG_t::kPi0, 0.f},
153+
{PDG_t::kKPlus, 1.f},
154+
{PDG_t::kK0Short, 1.f},
155+
{PDG_t::kK0Long, 1.f},
156+
{PDG_t::kProton, 2.f},
157+
{PDG_t::kNeutron, 2.f},
158+
{PDG_t::kSigmaPlus, 3.f},
159+
{PDG_t::kSigmaMinus, 3.f},
160+
{PDG_t::kLambda0, 3.f},
161+
{PDG_t::kSigma0, 3.f},
162+
{PDG_t::kXiMinus, 3.f},
163+
// TODO: potentially extend by xi0/eta/omega/rho/phi/Delta...
164+
};
165+
166+
if (auto iterMapPID = mapPID.find(absPDGCode); iterMapPID != mapPID.end()) {
167+
// LOGP(info, "scaling a {} with status code {} from process {}", particle.pdgCode(), particle.getGenStatusCode(), particle.getProcess());
168+
float pt = refParticle.pt();
169+
170+
// calculate particle fractions and corresponding weight for given reference particle
171+
std::vector<std::vector<float>> input = {{dNdEta}, {pt}, {iterMapPID->second}};
172+
float fracData = particleFractionsData.evalModel(input)[0];
173+
float fracMC = particleFractionsMC.evalModel(input)[0];
174+
if (fracMC) weight = fracData / fracMC;
175+
176+
if (enableQAHistos && std::abs(particle.eta()) < 0.8) {
177+
if (particle.isPhysicalPrimary()) {
178+
if (iterMapPID->first == PDG_t::kPiPlus) {
179+
histos.fill(HIST("frac/data/pion"), pt, fracData);
180+
histos.fill(HIST("frac/mc/pion"), pt, fracMC);
181+
histos.fill(HIST("weight/pion"), pt, weight);
182+
}
183+
if (iterMapPID->first == PDG_t::kKPlus) {
184+
histos.fill(HIST("frac/data/kaon"), pt, fracData);
185+
histos.fill(HIST("frac/mc/kaon"), pt, fracMC);
186+
histos.fill(HIST("weight/kaon"), pt, weight);
187+
}
188+
if (iterMapPID->first == PDG_t::kProton) {
189+
histos.fill(HIST("frac/data/proton"), pt, fracData);
190+
histos.fill(HIST("frac/mc/proton"), pt, fracMC);
191+
histos.fill(HIST("weight/proton"), pt, weight);
192+
}
193+
if (iterMapPID->first == PDG_t::kSigmaPlus || iterMapPID->first == PDG_t::kSigmaMinus) {
194+
histos.fill(HIST("frac/data/sigma"), pt, fracData);
195+
histos.fill(HIST("frac/mc/sigma"), pt, fracMC);
196+
histos.fill(HIST("weight/sigma"), pt, weight);
197+
}
198+
} else {
199+
// fill the decay daughter info only (usually there is max. one parent generation so no relevant double counting here)
200+
if (particle.getProcess() == TMCProcess::kPDecay) {
201+
histos.fill(HIST("weight/secDec"), pt, weight);
202+
} else {
203+
histos.fill(HIST("weight/secMat"), pt, weight);
204+
}
205+
}
206+
}
207+
}
208+
}
209+
}
210+
pccTable(weight);
211+
}
212+
}

0 commit comments

Comments
 (0)