Skip to content

Commit e167d13

Browse files
jezwilkinsonjezwilkinson
authored
beta-gamma cutoff, enable flags per species for NN (#5563)
1 parent 8f54982 commit e167d13

File tree

2 files changed

+55
-33
lines changed

2 files changed

+55
-33
lines changed

‎Common/TableProducer/PID/pidTPC.cxx‎

Lines changed: 28 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,16 @@ struct tpcPid {
102102
Configurable<int> pidTr{"pid-tr", -1, {"Produce PID information for the Triton mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};
103103
Configurable<int> pidHe{"pid-he", -1, {"Produce PID information for the Helium3 mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};
104104
Configurable<int> pidAl{"pid-al", -1, {"Produce PID information for the Alpha mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};
105+
Configurable<int> useNetworkEl{"useNetworkEl", 1, {"Switch for applying neural network on the electron mass hypothesis (if network enabled) (set to 0 to disable)"}};
106+
Configurable<int> useNetworkMu{"useNetworkMu", 1, {"Switch for applying neural network on the muon mass hypothesis (if network enabled) (set to 0 to disable)"}};
107+
Configurable<int> useNetworkPi{"useNetworkPi", 1, {"Switch for applying neural network on the pion mass hypothesis (if network enabled) (set to 0 to disable)"}};
108+
Configurable<int> useNetworkKa{"useNetworkKa", 1, {"Switch for applying neural network on the kaon mass hypothesis (if network enabled) (set to 0 to disable)"}};
109+
Configurable<int> useNetworkPr{"useNetworkPr", 1, {"Switch for applying neural network on the proton mass hypothesis (if network enabled) (set to 0 to disable)"}};
110+
Configurable<int> useNetworkDe{"useNetworkDe", 1, {"Switch for applying neural network on the deuteron mass hypothesis (if network enabled) (set to 0 to disable)"}};
111+
Configurable<int> useNetworkTr{"useNetworkTr", 1, {"Switch for applying neural network on the triton mass hypothesis (if network enabled) (set to 0 to disable)"}};
112+
Configurable<int> useNetworkHe{"useNetworkHe", 1, {"Switch for applying neural network on the helium3 mass hypothesis (if network enabled) (set to 0 to disable)"}};
113+
Configurable<int> useNetworkAl{"useNetworkAl", 1, {"Switch for applying neural network on the alpha mass hypothesis (if network enabled) (set to 0 to disable)"}};
114+
Configurable<float> networkBetaGammaCutoff{"networkBetaGammaCutoff", 0.45, {"Lower value of beta-gamma to override the NN application"}};
105115

106116
// Paramatrization configuration
107117
bool useCCDBParam = false;
@@ -202,11 +212,11 @@ struct tpcPid {
202212
}
203213
}
204214

205-
Partition<Trks> notTPCStandaloneTracks = ((aod::track::itsClusterSizes > (uint32_t)0) || (aod::track::trdPattern > (uint8_t)0) || (aod::track::tofExpMom > 0.f && aod::track::tofChi2 > 0.f)); // To count number of tracks for use in NN array
215+
Partition<Trks> notTPCStandaloneTracks = (aod::track::tpcNClsFindable > (uint8_t)0) && ((aod::track::itsClusterSizes > (uint32_t)0) || (aod::track::trdPattern > (uint8_t)0) || (aod::track::tofExpMom > 0.f && aod::track::tofChi2 > 0.f)); // To count number of tracks for use in NN array
206216
Partition<Trks> tracksWithTPC = (aod::track::tpcNClsFindable > (uint8_t)0);
207217

208218
void process(Coll const& collisions, Trks const& tracks,
209-
aod::BCsWithTimestamps const&)
219+
aod::BCsWithTimestamps const& bcs)
210220
{
211221

212222
const uint64_t outTable_size = tracks.size();
@@ -233,7 +243,7 @@ struct tpcPid {
233243
if (useNetworkCorrection) {
234244
auto start_network_total = std::chrono::high_resolution_clock::now();
235245
if (autofetchNetworks) {
236-
auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
246+
auto bc = bcs.begin();
237247
// Initialise correct TPC response object before NN setup (for NCl normalisation)
238248
if (useCCDBParam && ccdbTimestamp.value == 0 && !ccdb->isCachedObjectValid(ccdbPath.value, bc.timestamp())) { // Updating parametrisation only if the initial timestamp is 0
239249
if (recoPass.value == "") {
@@ -297,7 +307,7 @@ struct tpcPid {
297307
track_properties[counter_track_props + 1] = trk.tgl();
298308
track_properties[counter_track_props + 2] = trk.signed1Pt();
299309
track_properties[counter_track_props + 3] = o2::track::pid_constants::sMasses[i];
300-
track_properties[counter_track_props + 4] = collisions.iteratorAt(trk.collisionId()).multTPC() / 11000.;
310+
track_properties[counter_track_props + 4] = trk.has_collision() ? collisions.iteratorAt(trk.collisionId()).multTPC() / 11000. : 1.; // Dummy value in case no associated collision
301311
track_properties[counter_track_props + 5] = std::sqrt(nNclNormalization / trk.tpcNClsFound());
302312
counter_track_props += input_dimensions;
303313
}
@@ -346,7 +356,7 @@ struct tpcPid {
346356
}
347357
}
348358
// Check and fill enabled tables
349-
auto makeTable = [&trk, &collisions, &network_prediction, &count_tracks, &tracksForNet_size, this](const Configurable<int>& flag, auto& table, const o2::track::PID::ID pid) {
359+
auto makeTable = [&trk, &collisions, &network_prediction, &count_tracks, &tracksForNet_size, this](const Configurable<int>& flag, auto& table, const o2::track::PID::ID pid, bool speciesApplicationFlag) {
350360
if (flag.value != 1) {
351361
return;
352362
}
@@ -361,13 +371,14 @@ struct tpcPid {
361371
}
362372
}
363373
auto expSignal = response->GetExpectedSignal(trk, pid);
364-
auto expSigma = response->GetExpectedSigma(collisions.iteratorAt(trk.collisionId()), trk, pid);
365-
if (expSignal < 0. || expSigma < 0.) { // skip if expected signal invalid
374+
auto expSigma = trk.has_collision() ? response->GetExpectedSigma(collisions.iteratorAt(trk.collisionId()), trk, pid) : 0.07 * expSignal; // use default sigma value of 7% if no collision information to estimate resolution
375+
if (expSignal < 0. || expSigma < 0.) { // skip if expected signal invalid
366376
table(aod::pidtpc_tiny::binning::underflowBin);
367377
return;
368378
}
379+
float bg = trk.tpcInnerParam() / o2::track::pid_constants::sMasses[pid]; // estimated beta-gamma for network cutoff
369380

370-
if (useNetworkCorrection) {
381+
if (useNetworkCorrection && speciesApplicationFlag && trk.has_collision() && bg > networkBetaGammaCutoff) {
371382

372383
// Here comes the application of the network. The output--dimensions of the network dtermine the application: 1: mean, 2: sigma, 3: sigma asymmetric
373384
// For now only the option 2: sigma will be used. The other options are kept if there would be demand later on
@@ -389,15 +400,15 @@ struct tpcPid {
389400
}
390401
};
391402

392-
makeTable(pidEl, tablePIDEl, o2::track::PID::Electron);
393-
makeTable(pidMu, tablePIDMu, o2::track::PID::Muon);
394-
makeTable(pidPi, tablePIDPi, o2::track::PID::Pion);
395-
makeTable(pidKa, tablePIDKa, o2::track::PID::Kaon);
396-
makeTable(pidPr, tablePIDPr, o2::track::PID::Proton);
397-
makeTable(pidDe, tablePIDDe, o2::track::PID::Deuteron);
398-
makeTable(pidTr, tablePIDTr, o2::track::PID::Triton);
399-
makeTable(pidHe, tablePIDHe, o2::track::PID::Helium3);
400-
makeTable(pidAl, tablePIDAl, o2::track::PID::Alpha);
403+
makeTable(pidEl, tablePIDEl, o2::track::PID::Electron, useNetworkEl);
404+
makeTable(pidMu, tablePIDMu, o2::track::PID::Muon, useNetworkMu);
405+
makeTable(pidPi, tablePIDPi, o2::track::PID::Pion, useNetworkPi);
406+
makeTable(pidKa, tablePIDKa, o2::track::PID::Kaon, useNetworkKa);
407+
makeTable(pidPr, tablePIDPr, o2::track::PID::Proton, useNetworkPr);
408+
makeTable(pidDe, tablePIDDe, o2::track::PID::Deuteron, useNetworkDe);
409+
makeTable(pidTr, tablePIDTr, o2::track::PID::Triton, useNetworkTr);
410+
makeTable(pidHe, tablePIDHe, o2::track::PID::Helium3, useNetworkHe);
411+
makeTable(pidAl, tablePIDAl, o2::track::PID::Alpha, useNetworkAl);
401412

402413
if (trk.hasTPC() && (!skipTPCOnly || trk.hasITS() || trk.hasTRD() || trk.hasTOF())) {
403414
count_tracks++; // Increment network track counter only if (not skipping TPConly) or (is not TPConly)

‎Common/TableProducer/PID/pidTPCFull.cxx‎

Lines changed: 27 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,16 @@ struct tpcPidFull {
102102
Configurable<int> pidTr{"pid-tr", -1, {"Produce PID information for the Triton mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};
103103
Configurable<int> pidHe{"pid-he", -1, {"Produce PID information for the Helium3 mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};
104104
Configurable<int> pidAl{"pid-al", -1, {"Produce PID information for the Alpha mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};
105+
Configurable<int> useNetworkEl{"useNetworkEl", 1, {"Switch for applying neural network on the electron mass hypothesis (if network enabled) (set to 0 to disable)"}};
106+
Configurable<int> useNetworkMu{"useNetworkMu", 1, {"Switch for applying neural network on the muon mass hypothesis (if network enabled) (set to 0 to disable)"}};
107+
Configurable<int> useNetworkPi{"useNetworkPi", 1, {"Switch for applying neural network on the pion mass hypothesis (if network enabled) (set to 0 to disable)"}};
108+
Configurable<int> useNetworkKa{"useNetworkKa", 1, {"Switch for applying neural network on the kaon mass hypothesis (if network enabled) (set to 0 to disable)"}};
109+
Configurable<int> useNetworkPr{"useNetworkPr", 1, {"Switch for applying neural network on the proton mass hypothesis (if network enabled) (set to 0 to disable)"}};
110+
Configurable<int> useNetworkDe{"useNetworkDe", 1, {"Switch for applying neural network on the deuteron mass hypothesis (if network enabled) (set to 0 to disable)"}};
111+
Configurable<int> useNetworkTr{"useNetworkTr", 1, {"Switch for applying neural network on the triton mass hypothesis (if network enabled) (set to 0 to disable)"}};
112+
Configurable<int> useNetworkHe{"useNetworkHe", 1, {"Switch for applying neural network on the helium3 mass hypothesis (if network enabled) (set to 0 to disable)"}};
113+
Configurable<int> useNetworkAl{"useNetworkAl", 1, {"Switch for applying neural network on the alpha mass hypothesis (if network enabled) (set to 0 to disable)"}};
114+
Configurable<float> networkBetaGammaCutoff{"networkBetaGammaCutoff", 0.45, {"Lower value of beta-gamma to override the NN application"}};
105115

106116
// Paramatrization configuration
107117
bool useCCDBParam = false;
@@ -207,7 +217,7 @@ struct tpcPidFull {
207217
Partition<Trks> tracksWithTPC = (aod::track::tpcNClsFindable > (uint8_t)0);
208218

209219
void process(Coll const& collisions, Trks const& tracks,
210-
aod::BCsWithTimestamps const&)
220+
aod::BCsWithTimestamps const& bcs)
211221
{
212222

213223
const uint64_t outTable_size = tracks.size();
@@ -235,7 +245,7 @@ struct tpcPidFull {
235245

236246
auto start_network_total = std::chrono::high_resolution_clock::now();
237247
if (autofetchNetworks) {
238-
auto bc = collisions.iteratorAt(0).bc_as<aod::BCsWithTimestamps>();
248+
auto bc = bcs.begin();
239249
// Initialise correct TPC response object before NN setup (for NCl normalisation)
240250
if (useCCDBParam && ccdbTimestamp.value == 0 && !ccdb->isCachedObjectValid(ccdbPath.value, bc.timestamp())) { // Updating parametrisation only if the initial timestamp is 0
241251
if (recoPass.value == "") {
@@ -298,7 +308,7 @@ struct tpcPidFull {
298308
track_properties[counter_track_props + 1] = trk.tgl();
299309
track_properties[counter_track_props + 2] = trk.signed1Pt();
300310
track_properties[counter_track_props + 3] = o2::track::pid_constants::sMasses[i];
301-
track_properties[counter_track_props + 4] = collisions.iteratorAt(trk.collisionId()).multTPC() / 11000.;
311+
track_properties[counter_track_props + 4] = trk.has_collision() ? collisions.iteratorAt(trk.collisionId()).multTPC() / 11000. : 1.; // Dummy value in case no associated collision
302312
track_properties[counter_track_props + 5] = std::sqrt(nNclNormalization / trk.tpcNClsFound());
303313
counter_track_props += input_dimensions;
304314
}
@@ -347,7 +357,7 @@ struct tpcPidFull {
347357
}
348358
}
349359
// Check and fill enabled tables
350-
auto makeTable = [&trk, &collisions, &network_prediction, &count_tracks, &tracksForNet_size, this](const Configurable<int>& flag, auto& table, const o2::track::PID::ID pid) {
360+
auto makeTable = [&trk, &collisions, &network_prediction, &count_tracks, &tracksForNet_size, this](const Configurable<int>& flag, auto& table, const o2::track::PID::ID pid, bool speciesApplicationFlag) {
351361
if (flag.value != 1) {
352362
return;
353363
}
@@ -362,13 +372,14 @@ struct tpcPidFull {
362372
}
363373
}
364374
auto expSignal = response->GetExpectedSignal(trk, pid);
365-
auto expSigma = response->GetExpectedSigma(collisions.iteratorAt(trk.collisionId()), trk, pid);
366-
if (expSignal < 0. || expSigma < 0.) { // skip if expected signal invalid
375+
auto expSigma = trk.has_collision() ? response->GetExpectedSigma(collisions.iteratorAt(trk.collisionId()), trk, pid) : 0.07 * expSignal; // use default sigma value of 7% if no collision information to estimate resolution
376+
if (expSignal < 0. || expSigma < 0.) { // skip if expected signal invalid
367377
table(-999.f, -999.f);
368378
return;
369379
}
380+
float bg = trk.tpcInnerParam() / o2::track::pid_constants::sMasses[pid]; // estimated beta-gamma for network cutoff
370381

371-
if (useNetworkCorrection) {
382+
if (useNetworkCorrection && speciesApplicationFlag && trk.has_collision() && bg > networkBetaGammaCutoff) {
372383

373384
// Here comes the application of the network. The output--dimensions of the network dtermine the application: 1: mean, 2: sigma, 3: sigma asymmetric
374385
// For now only the option 2: sigma will be used. The other options are kept if there would be demand later on
@@ -395,15 +406,15 @@ struct tpcPidFull {
395406
}
396407
};
397408

398-
makeTable(pidEl, tablePIDEl, o2::track::PID::Electron);
399-
makeTable(pidMu, tablePIDMu, o2::track::PID::Muon);
400-
makeTable(pidPi, tablePIDPi, o2::track::PID::Pion);
401-
makeTable(pidKa, tablePIDKa, o2::track::PID::Kaon);
402-
makeTable(pidPr, tablePIDPr, o2::track::PID::Proton);
403-
makeTable(pidDe, tablePIDDe, o2::track::PID::Deuteron);
404-
makeTable(pidTr, tablePIDTr, o2::track::PID::Triton);
405-
makeTable(pidHe, tablePIDHe, o2::track::PID::Helium3);
406-
makeTable(pidAl, tablePIDAl, o2::track::PID::Alpha);
409+
makeTable(pidEl, tablePIDEl, o2::track::PID::Electron, useNetworkEl);
410+
makeTable(pidMu, tablePIDMu, o2::track::PID::Muon, useNetworkMu);
411+
makeTable(pidPi, tablePIDPi, o2::track::PID::Pion, useNetworkPi);
412+
makeTable(pidKa, tablePIDKa, o2::track::PID::Kaon, useNetworkKa);
413+
makeTable(pidPr, tablePIDPr, o2::track::PID::Proton, useNetworkPr);
414+
makeTable(pidDe, tablePIDDe, o2::track::PID::Deuteron, useNetworkDe);
415+
makeTable(pidTr, tablePIDTr, o2::track::PID::Triton, useNetworkTr);
416+
makeTable(pidHe, tablePIDHe, o2::track::PID::Helium3, useNetworkHe);
417+
makeTable(pidAl, tablePIDAl, o2::track::PID::Alpha, useNetworkAl);
407418

408419
if (trk.hasTPC() && (!skipTPCOnly || trk.hasITS() || trk.hasTRD() || trk.hasTOF())) {
409420
count_tracks++; // Increment network track counter only if track has TPC, and (not skipping TPConly) or (is not TPConly)

0 commit comments

Comments
 (0)