Skip to content
198 changes: 110 additions & 88 deletions Common/Tools/PID/pidTPCModule.h
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,15 @@
#include "Common/CCDB/ctpRateFetcher.h"
#include "Common/Core/CollisionTypeHelper.h"
#include "Common/Core/PID/TPCPIDResponse.h"
#include "Common/Core/RecoDecay.h"
#include "Common/Core/TableHelper.h"
#include "Common/DataModel/EventSelection.h"
#include "Common/DataModel/PIDResponseTPC.h"
#include "Common/TableProducer/PID/pidTPCBase.h" // IWYU pragma: keep
#include "Tools/ML/model.h"

#include <CommonConstants/MathConstants.h>
#include <CommonConstants/PhysicsConstants.h>
#include <DataFormatsParameters/GRPLHCIFData.h>
#include <Framework/AnalysisDataModel.h>
#include <Framework/AnalysisHelpers.h>
Expand All @@ -46,6 +49,7 @@
#include <TRandom.h>
#include <TString.h>

#include <array>
#include <chrono>
#include <cstddef>
#include <cstdint>
Expand All @@ -54,6 +58,7 @@
#include <memory>
#include <ratio>
#include <string>
#include <string_view>
#include <vector>

#include <math.h>
Expand Down Expand Up @@ -92,10 +97,10 @@

struct pidTPCConfigurables : o2::framework::ConfigurableGroup {
std::string prefix = "pidTPC";
o2::framework::Configurable<std::string> paramfile{"param-file", "", "Path to the parametrization object, if empty the parametrization is not taken from file"};

Check failure on line 100 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
o2::framework::Configurable<std::string> ccdbPath{"ccdbPath", "Analysis/PID/TPC/Response", "Path of the TPC parametrization on the CCDB"};
o2::framework::Configurable<std::string> recoPass{"recoPass", "", "Reconstruction pass name for CCDB query (automatically takes latest object for timestamp if blank)"};
o2::framework::Configurable<int64_t> ccdbTimestamp{"ccdb-timestamp", 0, "timestamp of the object used to query in CCDB the detector response. Exceptions: -1 gets the latest object, 0 gets the run dependent timestamp"};

Check failure on line 103 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
// Parameters for loading network from a file / downloading the file
o2::framework::Configurable<bool> useNetworkCorrection{"useNetworkCorrection", 0, "(bool) Wether or not to use the network correction for the TPC dE/dx signal"};
o2::framework::Configurable<bool> autofetchNetworks{"autofetchNetworks", 1, "(bool) Automatically fetches networks from CCDB for the correct run number"};
Expand All @@ -109,14 +114,14 @@
o2::framework::Configurable<int> savedEdxsCorrected{"savedEdxsCorrected", -1, {"Save table with corrected dE/dx calculated on the spot. 0: off, 1: on, -1: auto"}};
o2::framework::Configurable<bool> useCorrecteddEdx{"useCorrecteddEdx", false, "(bool) If true, use corrected dEdx value in Nsigma calculation instead of the one in the AO2D"};

o2::framework::Configurable<int> pidFullEl{"pid-full-el", -1, {"Produce PID information for the Electron mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};

Check failure on line 117 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
o2::framework::Configurable<int> pidFullMu{"pid-full-mu", -1, {"Produce PID information for the Muon mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};

Check failure on line 118 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
o2::framework::Configurable<int> pidFullPi{"pid-full-pi", -1, {"Produce PID information for the Pion mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};

Check failure on line 119 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
o2::framework::Configurable<int> pidFullKa{"pid-full-ka", -1, {"Produce PID information for the Kaon mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};

Check failure on line 120 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
o2::framework::Configurable<int> pidFullPr{"pid-full-pr", -1, {"Produce PID information for the Proton mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};

Check failure on line 121 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
o2::framework::Configurable<int> pidFullDe{"pid-full-de", -1, {"Produce PID information for the Deuterons mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};

Check failure on line 122 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
o2::framework::Configurable<int> pidFullTr{"pid-full-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)"}};

Check failure on line 123 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
o2::framework::Configurable<int> pidFullHe{"pid-full-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)"}};

Check failure on line 124 in Common/Tools/PID/pidTPCModule.h

View workflow job for this annotation

GitHub Actions / O2 linter

[name/configurable]

Use lowerCamelCase for names of configurables and use the same name for the struct member as for the JSON string. (Declare the type and names on the same line.)
o2::framework::Configurable<int> pidFullAl{"pid-full-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)"}};
o2::framework::Configurable<int> pidTinyEl{"pid-tiny-el", -1, {"Produce PID information for the Electron mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};
o2::framework::Configurable<int> pidTinyMu{"pid-tiny-mu", -1, {"Produce PID information for the Muon mass hypothesis, overrides the automatic setup: the corresponding table can be set off (0) or on (1)"}};
Expand All @@ -139,6 +144,9 @@
o2::framework::Configurable<int> useNetworkAl{"useNetworkAl", 1, {"Switch for applying neural network on the alpha mass hypothesis (if network enabled) (set to 0 to disable)"}};
o2::framework::Configurable<float> networkBetaGammaCutoff{"networkBetaGammaCutoff", 0.45, {"Lower value of beta-gamma to override the NN application"}};
o2::framework::Configurable<std::string> cfgPathGrpLhcIf{"ccdb-path-grplhcif", "GLO/Config/GRPLHCIF", "Path on the CCDB for the GRPLHCIF object"};

o2::framework::Configurable<float> phiEntranceCoeff1{"phiEntranceCoeff1", 1.026f, "phiEntrance = phi + phiEntranceCoeff1 * LightSpeedDm2PS * 0.5 * phiEntranceCoeff2 * 1/pT[GeV/c]"};
o2::framework::Configurable<float> phiEntranceCoeff2{"phiEntranceCoeff2", 85.f, "phiEntrance = phi + phiEntranceCoeff1 * LightSpeedDm2PS * 0.5 * phiEntranceCoeff2 * 1/pT[GeV/c]"};
};

// helper getter - FIXME should be separate
Expand Down Expand Up @@ -445,10 +453,49 @@
template <typename TCCDB, typename M, typename T, typename B>
std::vector<float> createNetworkPrediction(TCCDB& ccdb, soa::Join<aod::Collisions, aod::EvSels> const& collisions, M const& mults, T const& tracks, B const& bcs, const size_t size)
{
constexpr int NParticleTypes = 9;
constexpr double OneToKilo = 1.e-3;
constexpr int NanoToOne = 1000000000;
constexpr double MultiplicityNorm = 11000.;
constexpr double HadronicRateNormPp = 1500.;
constexpr double HadronicRateNormAa = 50.;
constexpr double Ft0cOccupancyNorm = 60000.;
constexpr float LightSpeedDm2PS = o2::constants::physics::LightSpeedCm2PS / 10.f;

struct NNVersionEntry {
std::string_view versionName{};
int numberOfFeatures{};
int versionNumber{};
};

std::vector<float> network_prediction;
constexpr std::array<NNVersionEntry, 6> nnVersionsDictionary{
{{"", 6, 1},
{"1", 6, 1},
{"2", 7, 2},
{"3", 8, 3},
{"4", 9, 4},
{"5", 9, 5}}};

enum IndexNnFeature : int {
IdxTpcInnerParam = 0,
IdxTgl,
IdxSigned1Pt,
IdxMass,
IdxMultiplicity,
IdxNClusters,
IdxFt0cOcc,
IdxHadronicRate,
IdxModPhi
};

auto start_network_total = std::chrono::high_resolution_clock::now();
constexpr int OldestNNVersionWithFt0c{2};
constexpr int OldestNNVersionWithHadronicRate{3};
constexpr int OldestNNVersionWithModPhi{4};
constexpr int NNVersionWithModPhiEntrance{5};

std::vector<float> networkPrediction;

const auto startNetworkTotal = std::chrono::high_resolution_clock::now();
if (pidTPCopts.autofetchNetworks) {
const auto& bc = bcs.begin();
// Initialise correct TPC response object before NN setup (for NCl normalisation)
Expand Down Expand Up @@ -486,7 +533,7 @@

if (bc.timestamp() < network.getValidityFrom() || bc.timestamp() > network.getValidityUntil()) { // fetches network only if the runnumbers change
LOG(info) << "Fetching network for timestamp: " << bc.timestamp();
bool retrieveSuccess = ccdb->getCCDBAccessor().retrieveBlob(pidTPCopts.networkPathCCDB.value, ".", metadata, bc.timestamp(), false, pidTPCopts.networkPathLocally.value, "", "", &headers);
const bool retrieveSuccess = ccdb->getCCDBAccessor().retrieveBlob(pidTPCopts.networkPathCCDB.value, ".", metadata, bc.timestamp(), false, pidTPCopts.networkPathLocally.value, "", "", &headers);
networkVersion = headers["NN-Version"];
if (retrieveSuccess) {
network.initModel(pidTPCopts.networkPathLocally.value, pidTPCopts.enableNetworkOptimizations.value, pidTPCopts.networkSetNumThreads.value, strtoul(headers["Valid-From"].c_str(), NULL, 0), strtoul(headers["Valid-Until"].c_str(), NULL, 0));
Expand All @@ -500,49 +547,48 @@
}

// Defining some network parameters
int input_dimensions = network.getNumInputNodes();
int output_dimensions = network.getNumOutputNodes();
const uint64_t track_prop_size = input_dimensions * size;
const uint64_t prediction_size = output_dimensions * size;
const int inputDimensions = network.getNumInputNodes();
const int outputDimensions = network.getNumOutputNodes();
const uint64_t trackPropSize = inputDimensions * size;
const uint64_t predictionSize = outputDimensions * size;

int nnVersion{0};
for (const auto& nnVersionEntry : nnVersionsDictionary) {
if (networkVersion == nnVersionEntry.versionName && inputDimensions == nnVersionEntry.numberOfFeatures) {
nnVersion = nnVersionEntry.versionNumber;
break;
}
}
if (nnVersion == 0) {
LOG(fatal) << "createNetworkPrediction(): networkVersion '" << networkVersion << "' and number of features " << inputDimensions << " are not compatible according to nnVersionsDictionary";
}

const int hadronicRateNorm = collsys == CollisionSystemType::kCollSyspp ? HadronicRateNormPp : HadronicRateNormAa;

network_prediction = std::vector<float>(prediction_size * 9); // For each mass hypotheses
networkPrediction = std::vector<float>(predictionSize * NParticleTypes); // For each mass hypotheses
const float nNclNormalization = response->GetNClNormalization();
float duration_network = 0;
float durationNetwork = 0;

std::vector<float> track_properties(track_prop_size);
uint64_t counter_track_props = 0;
int loop_counter = 0;
std::vector<float> trackProperties(trackPropSize);
uint64_t counterTrackProps = 0;
int loopCounter = 0;

// To load the Hadronic rate once for each collision
float hadronicRateBegin = 0.;
std::vector<float> hadronicRateForCollision(collisions.size(), 0.0f);
size_t i = 0;
size_t iCollision = 0;
for (const auto& collision : collisions) {
const auto& bc = collision.template bc_as<B>();
if (irSource.compare("") != 0) {
hadronicRateForCollision[i] = mRateFetcher.fetch(ccdb.service, bc.timestamp(), bc.runNumber(), irSource) * 1.e-3;
} else {
hadronicRateForCollision[i] = 0.0f;
hadronicRateForCollision[iCollision] = mRateFetcher.fetch(ccdb.service, bc.timestamp(), bc.runNumber(), irSource) * OneToKilo;
}
i++;
}
auto bc = bcs.begin();
if (irSource.compare("") != 0) {
hadronicRateBegin = mRateFetcher.fetch(ccdb.service, bc.timestamp(), bc.runNumber(), irSource) * 1.e-3; // kHz
} else {
hadronicRateBegin = 0.0f;
++iCollision;
}
const auto bc = bcs.begin();
const float hadronicRateBegin = irSource.compare("") != 0 ? mRateFetcher.fetch(ccdb.service, bc.timestamp(), bc.runNumber(), irSource) * OneToKilo : 0.f;

// Filling a std::vector<float> to be evaluated by the network
// Evaluation on single tracks brings huge overhead: Thus evaluation is done on one large vector
static constexpr int NParticleTypes = 9;
constexpr int ExpectedInputDimensionsNNV2 = 7;
constexpr int ExpectedInputDimensionsNNV3 = 8;
constexpr int ExpectedInputDimensionsNNV4 = 9;
constexpr auto NetworkVersionV2 = "2";
constexpr auto NetworkVersionV3 = "3";
constexpr auto NetworkVersionV4 = "4";
for (int j = 0; j < NParticleTypes; j++) { // Loop over particle number for which network correction is used
for (int jParticleType = 0; jParticleType < NParticleTypes; ++jParticleType) { // Loop over particle number for which network correction is used
for (auto const& trk : tracks) {
if (!trk.hasTPC()) {
continue;
Expand All @@ -552,74 +598,50 @@
continue;
}
}
track_properties[counter_track_props] = trk.tpcInnerParam();
track_properties[counter_track_props + 1] = trk.tgl();
track_properties[counter_track_props + 2] = trk.signed1Pt();
track_properties[counter_track_props + 3] = o2::track::pid_constants::sMasses[j];
track_properties[counter_track_props + 4] = (trk.has_collision() && mults.size() > 0) ? mults[trk.collisionId()] / 11000. : 1.;
track_properties[counter_track_props + 5] = std::sqrt(nNclNormalization / trk.tpcNClsFound());
if (input_dimensions == ExpectedInputDimensionsNNV2 && networkVersion == NetworkVersionV2) {
track_properties[counter_track_props + 6] = (trk.has_collision() && mults.size() > 0) ? collisions.iteratorAt(trk.collisionId()).ft0cOccupancyInTimeRange() / 60000. : 1.;
const bool isGoodTrack = trk.has_collision() && mults.size() > 0;
trackProperties[counterTrackProps + IdxTpcInnerParam] = trk.tpcInnerParam();
trackProperties[counterTrackProps + IdxTgl] = trk.tgl();
trackProperties[counterTrackProps + IdxSigned1Pt] = trk.signed1Pt();
trackProperties[counterTrackProps + IdxMass] = o2::track::pid_constants::sMasses[jParticleType];
trackProperties[counterTrackProps + IdxMultiplicity] = isGoodTrack ? mults[trk.collisionId()] / MultiplicityNorm : 1.;
trackProperties[counterTrackProps + IdxNClusters] = std::sqrt(nNclNormalization / trk.tpcNClsFound());
if (nnVersion >= OldestNNVersionWithFt0c) {
trackProperties[counterTrackProps + IdxFt0cOcc] = isGoodTrack ? collisions.iteratorAt(trk.collisionId()).ft0cOccupancyInTimeRange() / Ft0cOccupancyNorm : 1.;
}
if (input_dimensions == ExpectedInputDimensionsNNV3 && networkVersion == NetworkVersionV3) {
track_properties[counter_track_props + 6] = (trk.has_collision() && mults.size() > 0) ? collisions.iteratorAt(trk.collisionId()).ft0cOccupancyInTimeRange() / 60000. : 1.;
if (trk.has_collision() && mults.size() > 0) {
if (collsys == CollisionSystemType::kCollSyspp) {
track_properties[counter_track_props + 7] = hadronicRateForCollision[trk.collisionId()] / 1500.;
} else {
track_properties[counter_track_props + 7] = hadronicRateForCollision[trk.collisionId()] / 50.;
}
} else {
// asign Hadronic Rate at beginning of run if track does not belong to a collision
if (collsys == CollisionSystemType::kCollSyspp) {
track_properties[counter_track_props + 7] = hadronicRateBegin / 1500.;
} else {
track_properties[counter_track_props + 7] = hadronicRateBegin / 50.;
}
}
if (nnVersion >= OldestNNVersionWithHadronicRate) {
const float hadronicRate = isGoodTrack ? hadronicRateForCollision[trk.collisionId()] : hadronicRateBegin;
trackProperties[counterTrackProps + IdxHadronicRate] = hadronicRate / hadronicRateNorm;
}

if (input_dimensions == ExpectedInputDimensionsNNV4 && networkVersion == NetworkVersionV4) {
track_properties[counter_track_props + 6] = (trk.has_collision() && mults.size() > 0) ? collisions.iteratorAt(trk.collisionId()).ft0cOccupancyInTimeRange() / 60000. : 1.;
if (trk.has_collision() && mults.size() > 0) {
if (collsys == CollisionSystemType::kCollSyspp) {
track_properties[counter_track_props + 7] = hadronicRateForCollision[trk.collisionId()] / 1500.;
} else {
track_properties[counter_track_props + 7] = hadronicRateForCollision[trk.collisionId()] / 50.;
}
} else {
// asign Hadronic Rate at beginning of run if track does not belong to a collision
if (collsys == CollisionSystemType::kCollSyspp) {
track_properties[counter_track_props + 7] = hadronicRateBegin / 1500.;
} else {
track_properties[counter_track_props + 7] = hadronicRateBegin / 50.;
}
if (nnVersion >= OldestNNVersionWithModPhi) {
float phi = trk.phi();
if (nnVersion == NNVersionWithModPhiEntrance) {
phi += pidTPCopts.phiEntranceCoeff1 * LightSpeedDm2PS * 0.5 * pidTPCopts.phiEntranceCoeff2 * trk.signed1Pt();
}
track_properties[counter_track_props + 8] = std::fmod(std::fmod(trk.phi(), 2 * M_PI) + 2 * M_PI, M_PI / 9.0);
trackProperties[counterTrackProps + IdxModPhi] = RecoDecay::constrainAngle(phi, 0.f, o2::constants::math::NSectors);
}
counter_track_props += input_dimensions;
counterTrackProps += inputDimensions;
}

auto start_network_eval = std::chrono::high_resolution_clock::now();
float* output_network = network.evalModel(track_properties);
auto stop_network_eval = std::chrono::high_resolution_clock::now();
duration_network += std::chrono::duration<float, std::ratio<1, 1000000000>>(stop_network_eval - start_network_eval).count();
for (uint64_t k = 0; k < prediction_size; k += output_dimensions) {
for (int l = 0; l < output_dimensions; l++) {
network_prediction[k + l + prediction_size * loop_counter] = output_network[k + l];
const auto startNetworkEval = std::chrono::high_resolution_clock::now();
const float* const outputNetwork = network.evalModel(trackProperties);
const auto stopNetworkEval = std::chrono::high_resolution_clock::now();
durationNetwork += std::chrono::duration<float, std::ratio<1, NanoToOne>>(stopNetworkEval - startNetworkEval).count();
for (uint64_t kPrediction = 0; kPrediction < predictionSize; kPrediction += outputDimensions) {
for (int lOutputDim = 0; lOutputDim < outputDimensions; ++lOutputDim) {
networkPrediction[kPrediction + lOutputDim + predictionSize * loopCounter] = outputNetwork[kPrediction + lOutputDim];
}
}

counter_track_props = 0;
loop_counter += 1;
counterTrackProps = 0;
++loopCounter;
}
track_properties.clear();
trackProperties.clear();

auto stop_network_total = std::chrono::high_resolution_clock::now();
LOG(debug) << "Neural Network for the TPC PID response correction: Time per track (eval ONNX): " << duration_network / (size * 9) << "ns ; Total time (eval ONNX): " << duration_network / 1000000000 << " s";
LOG(debug) << "Neural Network for the TPC PID response correction: Time per track (eval + overhead): " << std::chrono::duration<float, std::ratio<1, 1000000000>>(stop_network_total - start_network_total).count() / (size * 9) << "ns ; Total time (eval + overhead): " << std::chrono::duration<float, std::ratio<1, 1000000000>>(stop_network_total - start_network_total).count() / 1000000000 << " s";
const auto stopNetworkTotal = std::chrono::high_resolution_clock::now();
LOG(debug) << "Neural Network for the TPC PID response correction: Time per track (eval ONNX): " << durationNetwork / (size * NParticleTypes) << "ns ; Total time (eval ONNX): " << durationNetwork / NanoToOne << " s";
LOG(debug) << "Neural Network for the TPC PID response correction: Time per track (eval + overhead): " << std::chrono::duration<float, std::ratio<1, NanoToOne>>(stopNetworkTotal - startNetworkTotal).count() / (size * NParticleTypes) << "ns ; Total time (eval + overhead): " << std::chrono::duration<float, std::ratio<1, NanoToOne>>(stopNetworkTotal - startNetworkTotal).count() / NanoToOne << " s";

return network_prediction;
return networkPrediction;
}

//__________________________________________________
Expand Down
Loading