Skip to content

Commit 0754513

Browse files
scannitoscannito
authored
[PWGLF] New analysis method + resonanceTreeCreator.cxx (#13056)
1 parent 55a59ee commit 0754513

File tree

3 files changed

+750
-333
lines changed

3 files changed

+750
-333
lines changed

‎PWGLF/TableProducer/Resonances/CMakeLists.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,3 +54,8 @@ o2physics_add_dpl_workflow(cksspinalignment
5454
SOURCES cksspinalignment.cxx
5555
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
5656
COMPONENT_NAME Analysis)
57+
58+
o2physics_add_dpl_workflow(resonance-tree-creator
59+
SOURCES resonanceTreeCreator.cxx
60+
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
61+
COMPONENT_NAME Analysis)

‎PWGLF/TableProducer/Resonances/resonanceTreeCreator.cxx‎

Lines changed: 344 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,344 @@
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 resonanceTreeCreator.cxx
13+
/// \brief Produces a TTree with machine learning variables for resonances in the LF group
14+
/// \author Stefano Cannito (stefano.cannito@cern.ch)
15+
16+
#include "PWGLF/DataModel/mcCentrality.h"
17+
#include "PWGLF/Utils/inelGt.h"
18+
19+
#include "Common/Core/RecoDecay.h"
20+
#include "Common/Core/TrackSelection.h"
21+
#include "Common/Core/trackUtilities.h"
22+
#include "Common/DataModel/Centrality.h"
23+
#include "Common/DataModel/EventSelection.h"
24+
#include "Common/DataModel/Multiplicity.h"
25+
#include "Common/DataModel/PIDResponse.h"
26+
#include "Common/DataModel/TrackSelectionTables.h"
27+
28+
#include "CommonConstants/PhysicsConstants.h"
29+
#include "Framework/ASoAHelpers.h"
30+
#include "Framework/AnalysisDataModel.h"
31+
#include "Framework/AnalysisTask.h"
32+
#include "Framework/HistogramRegistry.h"
33+
#include "Framework/O2DatabasePDGPlugin.h"
34+
#include "Framework/runDataProcessing.h"
35+
#include "ReconstructionDataFormats/Track.h"
36+
37+
#include <Math/Vector4D.h>
38+
#include <TPDGCode.h>
39+
40+
#include <algorithm>
41+
#include <array>
42+
#include <cmath>
43+
#include <cstdlib>
44+
#include <string>
45+
#include <utility>
46+
#include <vector>
47+
48+
using namespace o2;
49+
using namespace o2::framework;
50+
using namespace o2::framework::expressions;
51+
52+
namespace o2::aod
53+
{
54+
namespace resomlcandidates
55+
{
56+
// Multiplicity class
57+
DECLARE_SOA_COLUMN(MultClass, multClass, float); //! Event multiplicity class
58+
// Daughter 1
59+
DECLARE_SOA_COLUMN(PtDaughter1, ptdaughter1, float); //! Transverse momentum of daughter1 (GeV/c)
60+
DECLARE_SOA_COLUMN(PDaughter1, pdaughter1, float); //! Momentum of daughter1 (GeV/c)
61+
DECLARE_SOA_COLUMN(PhiDaughter1, phiDaughter1, float); //! Azimuthal angle of daughter1 (rad)
62+
DECLARE_SOA_COLUMN(EtaDaughter1, etaDaughter1, float); //! Pseudorapidity of daughter1
63+
DECLARE_SOA_COLUMN(YDaughter1, yDaughter1, float); //! Rapidity of daughter1
64+
DECLARE_SOA_COLUMN(DCAxyDaughter1, dcaDaughter1, float); //! DCA of daughter1 to primary vertex (cm)
65+
DECLARE_SOA_COLUMN(DCAzDaughter1, dcaZDaughter1, float); //! DCA of daughter1 to primary vertex in z (cm)
66+
DECLARE_SOA_COLUMN(NSigTpcPi1, nSigTpcPi1, float); //! TPC Nsigma separation for daughter1 with pion mass hypothesis
67+
DECLARE_SOA_COLUMN(NSigTpcKa1, nSigTpcKa1, float); //! TPC Nsigma separation for daughter1 with kaon mass hypothesis
68+
DECLARE_SOA_COLUMN(NSigTofPi1, nSigTofPi1, float); //! TOF Nsigma separation for daughter1 with pion mass hypothesis
69+
DECLARE_SOA_COLUMN(NSigTofKa1, nSigTofKa1, float); //! TOF Nsigma separation for daughter1 with kaon mass hypothesis
70+
DECLARE_SOA_COLUMN(NSigTpcTofPi1, nSigTpcTofPi1, float); //! TPC and TOF combined Nsigma separation for daughter1 with pion mass hypothesis
71+
DECLARE_SOA_COLUMN(NSigTpcTofKa1, nSigTpcTofKa1, float); //! TPC and TOF combined Nsigma separation for daughter1 with kaon mass hypothesis
72+
// Daughter 2
73+
DECLARE_SOA_COLUMN(PtDaughter2, ptdaughter2, float); //! Transverse momentum of daughter2 (GeV/c)
74+
DECLARE_SOA_COLUMN(PDaughter2, pdaughter2, float); //! Momentum of daughter2 (in GeV/c)
75+
DECLARE_SOA_COLUMN(PhiDaughter2, phiDaughter2, float); //! Azimuthal angle of daughter2 (rad)
76+
DECLARE_SOA_COLUMN(EtaDaughter2, etaDaughter2, float); //! Pseudorapidity of daughter2
77+
DECLARE_SOA_COLUMN(YDaughter2, yDaughter2, float); //! Rapidity of daughter2
78+
DECLARE_SOA_COLUMN(DCAxyDaughter2, dcaDaughter2, float); //! DCA of daughter2 to primary vertex (cm)
79+
DECLARE_SOA_COLUMN(DCAzDaughter2, dcaZDaughter2, float); //! DCA of daughter2 to primary vertex in z (cm)
80+
DECLARE_SOA_COLUMN(NSigTpcPi2, nSigTpcPi2, float); //! TPC Nsigma separation for daughter2 with pion mass hypothesis
81+
DECLARE_SOA_COLUMN(NSigTpcKa2, nSigTpcKa2, float); //! TPC Nsigma separation for daughter2 with kaon mass hypothesis
82+
DECLARE_SOA_COLUMN(NSigTofPi2, nSigTofPi2, float); //! TOF Nsigma separation for daughter2 with pion mass hypothesis
83+
DECLARE_SOA_COLUMN(NSigTofKa2, nSigTofKa2, float); //! TOF Nsigma separation for daughter2 with kaon mass hypothesis
84+
DECLARE_SOA_COLUMN(NSigTpcTofPi2, nSigTpcTofPi2, float); //! TPC and TOF combined Nsigma separation for daughter2 with pion mass hypothesis
85+
DECLARE_SOA_COLUMN(NSigTpcTofKa2, nSigTpcTofKa2, float); //! TPC and TOF combined Nsigma separation for daughter2 with kaon mass hypothesis
86+
// Candidate
87+
DECLARE_SOA_COLUMN(M, m, float); //! Invariant mass of candidate (GeV/c2)
88+
DECLARE_SOA_COLUMN(Pt, pt, float); //! Transverse momentum of candidate (GeV/c)
89+
DECLARE_SOA_COLUMN(P, p, float); //! Momentum of candidate (GeV/c)
90+
DECLARE_SOA_COLUMN(Phi, phi, float); //! Azimuth angle of candidate
91+
DECLARE_SOA_COLUMN(Eta, eta, float); //! Pseudorapidity of candidate
92+
DECLARE_SOA_COLUMN(Y, y, float); //! Rapidity of candidate
93+
DECLARE_SOA_COLUMN(Sign, sign, int8_t); //! Sign of the candidate
94+
DECLARE_SOA_COLUMN(IsTruePhi, isTruePhi, bool); //! Flag to indicate if the candidate is a phi meson
95+
} // namespace resomlcandidates
96+
97+
DECLARE_SOA_TABLE(ResoCandidates, "AOD", "RESOCANDIDATES",
98+
resomlcandidates::M,
99+
resomlcandidates::Pt,
100+
resomlcandidates::P,
101+
resomlcandidates::Phi,
102+
resomlcandidates::Eta,
103+
resomlcandidates::Y,
104+
resomlcandidates::Sign,
105+
resomlcandidates::IsTruePhi);
106+
107+
DECLARE_SOA_TABLE(ResoMLCandidates, "AOD", "RESOMLCANDIDATES",
108+
resomlcandidates::MultClass,
109+
resomlcandidates::PtDaughter1,
110+
resomlcandidates::PDaughter1,
111+
resomlcandidates::PhiDaughter1,
112+
resomlcandidates::EtaDaughter1,
113+
resomlcandidates::YDaughter1,
114+
resomlcandidates::DCAxyDaughter1,
115+
resomlcandidates::DCAzDaughter1,
116+
resomlcandidates::NSigTpcPi1,
117+
resomlcandidates::NSigTpcKa1,
118+
resomlcandidates::NSigTofPi1,
119+
resomlcandidates::NSigTofKa1,
120+
resomlcandidates::NSigTpcTofPi1,
121+
resomlcandidates::NSigTpcTofKa1,
122+
resomlcandidates::PtDaughter2,
123+
resomlcandidates::PDaughter2,
124+
resomlcandidates::PhiDaughter2,
125+
resomlcandidates::EtaDaughter2,
126+
resomlcandidates::YDaughter2,
127+
resomlcandidates::DCAxyDaughter2,
128+
resomlcandidates::DCAzDaughter2,
129+
resomlcandidates::NSigTpcPi2,
130+
resomlcandidates::NSigTpcKa2,
131+
resomlcandidates::NSigTofPi2,
132+
resomlcandidates::NSigTofKa2,
133+
resomlcandidates::NSigTpcTofPi2,
134+
resomlcandidates::NSigTpcTofKa2,
135+
resomlcandidates::M,
136+
resomlcandidates::Pt,
137+
resomlcandidates::P,
138+
resomlcandidates::Phi,
139+
resomlcandidates::Eta,
140+
resomlcandidates::Y,
141+
resomlcandidates::Sign);
142+
143+
namespace resomlselection
144+
{
145+
DECLARE_SOA_COLUMN(PhiBDTScore, gammaBDTScore, float);
146+
} // namespace resomlselection
147+
148+
DECLARE_SOA_TABLE(ResoPhiMLSelection, "AOD", "RESOPHIMLSELECTION",
149+
resomlselection::PhiBDTScore);
150+
} // namespace o2::aod
151+
152+
struct resonanceTreeCreator {
153+
// Production of the TTree
154+
Produces<aod::ResoCandidates> resoCandidates;
155+
Produces<aod::ResoMLCandidates> resoMLCandidates;
156+
157+
// Configurables for track selection
158+
Configurable<float> cfgCutCharge{"cfgCutCharge", 0.0f, "Cut on the signed transverse momentum to select positive or negative tracks"};
159+
160+
// Configurables for production of training samples
161+
Configurable<bool> fillOnlySignal{"fillOnlySignal", false, "Flag to fill derived tables with signal for ML trainings"};
162+
Configurable<bool> fillOnlyBackground{"fillOnlyBackground", false, "Flag to fill derived tables with background for ML trainings"};
163+
Configurable<float> downSampleBkgFactor{"downSampleBkgFactor", 0.5f, "Fraction of background candidates to keep for ML trainings"};
164+
Configurable<float> ptMaxForDownSample{"ptMaxForDownSample", 10.0f, "Maximum pt for the application of the downsampling factor"};
165+
166+
// Defining the type of the collisions for data and MC
167+
using SelCollisions = soa::Join<aod::Collisions, aod::EvSels, aod::CentFT0Ms, aod::PVMults>;
168+
using SimCollisions = soa::Join<SelCollisions, aod::McCollisionLabels>;
169+
170+
// Defining the type of the tracks for data and MC
171+
using FullTracks = soa::Join<aod::Tracks, aod::TracksExtra, aod::TracksDCA, aod::TrackSelection, aod::pidTPCFullPi,
172+
aod::pidTPCFullKa, aod::pidTOFFullPi, aod::pidTOFFullKa, aod::pidTPCFullPr, aod::pidTOFFullPr>;
173+
using FullMCTracks = soa::Join<FullTracks, aod::McTrackLabels>;
174+
175+
Partition<FullTracks> posTracks = aod::track::signed1Pt > cfgCutCharge;
176+
Partition<FullTracks> negTracks = aod::track::signed1Pt < cfgCutCharge;
177+
178+
Partition<FullMCTracks> posMCTracks = aod::track::signed1Pt > cfgCutCharge;
179+
Partition<FullMCTracks> negMCTracks = aod::track::signed1Pt < cfgCutCharge;
180+
181+
Preslice<aod::Tracks> perColl = aod::track::collisionId;
182+
Preslice<aod::McParticles> perMCColl = aod::mcparticle::mcCollisionId;
183+
184+
// Necessary to flag INEL>0 events in GenMC
185+
Service<o2::framework::O2DatabasePDG> pdgDB;
186+
187+
// Cache for manual slicing
188+
SliceCache cache;
189+
190+
enum ParticleType {
191+
Pi,
192+
Ka,
193+
Pr
194+
};
195+
196+
// Constants
197+
double massPi = o2::constants::physics::MassPiPlus;
198+
double massKa = o2::constants::physics::MassKPlus;
199+
200+
void init(InitContext&)
201+
{
202+
}
203+
204+
// Combine Nsigma values from TPC and TOF
205+
template <int massHypo, typename T>
206+
float combineNSigma(const T& track)
207+
{
208+
float nSigmaTPC, nSigmaTOF;
209+
switch (massHypo) {
210+
case Pi:
211+
nSigmaTPC = track.tpcNSigmaPi();
212+
nSigmaTOF = track.tofNSigmaPi();
213+
break;
214+
case Ka:
215+
nSigmaTPC = track.tpcNSigmaKa();
216+
nSigmaTOF = track.tofNSigmaKa();
217+
break;
218+
case Pr:
219+
nSigmaTPC = track.tpcNSigmaPr();
220+
nSigmaTOF = track.tofNSigmaPr();
221+
break;
222+
default:
223+
break;
224+
}
225+
226+
static constexpr float defaultNSigmaTolerance = .1f;
227+
static constexpr float defaultNSigma = -999.f + defaultNSigmaTolerance;
228+
229+
if (nSigmaTPC > defaultNSigma && nSigmaTOF > defaultNSigma)
230+
return std::sqrt(0.5f * std::pow(nSigmaTPC, 2) + std::pow(nSigmaTOF, 2));
231+
if (nSigmaTPC > defaultNSigma)
232+
return std::abs(nSigmaTPC);
233+
if (nSigmaTOF > defaultNSigma)
234+
return std::abs(nSigmaTOF);
235+
return nSigmaTPC;
236+
}
237+
238+
// Reconstruct the candidate 4-momentum from two daughter tracks
239+
template <typename T>
240+
ROOT::Math::PxPyPzMVector recMother(const T& track1, const T& track2, float masstrack1, float masstrack2)
241+
{
242+
ROOT::Math::PxPyPzMVector daughter1(track1.px(), track1.py(), track1.pz(), masstrack1); // set the daughter1 4-momentum
243+
ROOT::Math::PxPyPzMVector daughter2(track2.px(), track2.py(), track2.pz(), masstrack2); // set the daughter2 4-momentum
244+
ROOT::Math::PxPyPzMVector mother = daughter1 + daughter2; // calculate the mother 4-momentum
245+
246+
return mother;
247+
}
248+
249+
template <typename T1, typename T2>
250+
void fillCandidateTree4ML(const T1& collision, const T2& track1, const T2& track2, float masstrack1, float masstrack2
251+
/*std::optional<std::reference_wrapper<const aod::McParticles>> mcParticles = std::nullopt*/)
252+
{
253+
auto tpctofPi1 = combineNSigma<Pi>(track1);
254+
auto tpctofKa1 = combineNSigma<Ka>(track1);
255+
auto tpctofPi2 = combineNSigma<Pi>(track2);
256+
auto tpctofKa2 = combineNSigma<Ka>(track2);
257+
258+
ROOT::Math::PxPyPzMVector recCandidate = recMother(track1, track2, masstrack1, masstrack2);
259+
260+
if (downSampleBkgFactor < 1.) {
261+
float pseudoRndm = track1.pt() * 1000. - static_cast<int64_t>(track1.pt() * 1000);
262+
if (recCandidate.pt() < ptMaxForDownSample && pseudoRndm >= downSampleBkgFactor)
263+
return;
264+
}
265+
266+
resoMLCandidates(collision.centFT0M(),
267+
track1.pt(), track1.p(), track1.phi(), track1.eta(), track1.rapidity(masstrack1), track1.dcaXY(), track1.dcaZ(),
268+
track1.tpcNSigmaPi(), track1.tpcNSigmaKa(), track1.tofNSigmaPi(), track1.tofNSigmaKa(), tpctofPi1, tpctofKa1,
269+
track2.pt(), track2.p(), track2.phi(), track2.eta(), track2.rapidity(masstrack2), track2.dcaXY(), track2.dcaZ(),
270+
track2.tpcNSigmaPi(), track2.tpcNSigmaKa(), track2.tofNSigmaPi(), track2.tofNSigmaKa(), tpctofPi2, tpctofKa2,
271+
recCandidate.M(), recCandidate.Pt(), recCandidate.P(), recCandidate.Phi(),
272+
recCandidate.Eta(), recCandidate.Rapidity(), track1.sign() + track2.sign());
273+
}
274+
275+
template <typename T>
276+
bool isMCPhi(const T& track1, const T& track2, const aod::McParticles& mcParticles)
277+
{
278+
if (!track1.has_mcParticle() || !track2.has_mcParticle())
279+
return false; // Skip filling if no MC truth is available for both tracks
280+
281+
auto mcTrack1 = mcParticles.rawIteratorAt(track1.mcParticleId());
282+
auto mcTrack2 = mcParticles.rawIteratorAt(track2.mcParticleId());
283+
284+
if (mcTrack1.pdgCode() != PDG_t::kKPlus || !mcTrack1.isPhysicalPrimary())
285+
return false; // Skip filling if the first track is not a primary K+
286+
if (mcTrack2.pdgCode() != PDG_t::kKMinus || !mcTrack2.isPhysicalPrimary())
287+
return false; // Skip filling if the second track is not a primary K-
288+
289+
const auto mcTrack1MotherIndexes = mcTrack1.mothersIds();
290+
const auto mcTrack2MotherIndexes = mcTrack2.mothersIds();
291+
292+
for (const auto& mcTrack1MotherIndex : mcTrack1MotherIndexes) {
293+
for (const auto& mcTrack2MotherIndex : mcTrack2MotherIndexes) {
294+
if (mcTrack1MotherIndex != mcTrack2MotherIndex)
295+
continue;
296+
297+
const auto mother = mcParticles.rawIteratorAt(mcTrack1MotherIndex);
298+
if (mother.pdgCode() == o2::constants::physics::Pdg::kPhi)
299+
return true;
300+
}
301+
}
302+
return false;
303+
}
304+
305+
void processData4ML(SelCollisions::iterator const& collision, FullTracks const&)
306+
{
307+
auto posThisColl = posTracks->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache);
308+
auto negThisColl = negTracks->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache);
309+
310+
for (const auto& track1 : posThisColl) {
311+
for (const auto& track2 : negThisColl) {
312+
// Fill the ResoMLCandidates table with candidates in Data
313+
fillCandidateTree4ML(collision, track1, track2, massKa, massKa);
314+
}
315+
}
316+
}
317+
318+
PROCESS_SWITCH(resonanceTreeCreator, processData4ML, "Fill ResoMLCandidates in Data", true);
319+
320+
void processMC4ML(SimCollisions::iterator const& collision, FullMCTracks const&, aod::McParticles const& mcParticles)
321+
{
322+
auto posThisColl = posMCTracks->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache);
323+
auto negThisColl = negMCTracks->sliceByCached(aod::track::collisionId, collision.globalIndex(), cache);
324+
325+
for (const auto& track1 : posThisColl) {
326+
for (const auto& track2 : negThisColl) {
327+
if (fillOnlySignal && !isMCPhi(track1, track2, mcParticles))
328+
return; // Skip filling if only signal is requested and not a phi in MC truth
329+
if (fillOnlyBackground && isMCPhi(track1, track2, mcParticles))
330+
return; // Skip filling if only background is requested and a phi in MC truth
331+
332+
// Fill the ResoMLCandidates table with candidates in MC
333+
fillCandidateTree4ML(collision, track1, track2, massKa, massKa);
334+
}
335+
}
336+
}
337+
338+
PROCESS_SWITCH(resonanceTreeCreator, processMC4ML, "Fill ResoMLCandidates in MC", false);
339+
};
340+
341+
WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
342+
{
343+
return WorkflowSpec{adaptAnalysisTask<resonanceTreeCreator>(cfgc)};
344+
}

0 commit comments

Comments
 (0)