Skip to content

Commit 40344b5

Browse files
scannitoscannito
committed
First commit for ResonanceTreeCreator
1 parent 1b3fd6f commit 40344b5

File tree

2 files changed

+339
-0
lines changed

2 files changed

+339
-0
lines changed

‎PWGLF/TableProducer/Resonances/CMakeLists.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,8 @@ o2physics_add_dpl_workflow(heptaquarktable
4949
SOURCES HeptaQuarktable.cxx
5050
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DetectorsVertexing
5151
COMPONENT_NAME Analysis)
52+
53+
o2physics_add_dpl_workflow(resonance-tree-creator
54+
SOURCES resonanceTreeCreator.cxx
55+
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DetectorsBase
56+
COMPONENT_NAME Analysis)

‎PWGLF/TableProducer/Resonances/resonanceTreeCreator.cxx‎

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

0 commit comments

Comments
 (0)