diff --git a/Common/Tools/PID/pidTPCModule.h b/Common/Tools/PID/pidTPCModule.h index 9e8d93c9e43..78402587660 100644 --- a/Common/Tools/PID/pidTPCModule.h +++ b/Common/Tools/PID/pidTPCModule.h @@ -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 +#include #include #include #include @@ -46,6 +49,7 @@ #include #include +#include #include #include #include @@ -54,6 +58,7 @@ #include #include #include +#include #include #include @@ -139,6 +144,9 @@ struct pidTPCConfigurables : o2::framework::ConfigurableGroup { o2::framework::Configurable useNetworkAl{"useNetworkAl", 1, {"Switch for applying neural network on the alpha mass hypothesis (if network enabled) (set to 0 to disable)"}}; o2::framework::Configurable networkBetaGammaCutoff{"networkBetaGammaCutoff", 0.45, {"Lower value of beta-gamma to override the NN application"}}; o2::framework::Configurable cfgPathGrpLhcIf{"ccdb-path-grplhcif", "GLO/Config/GRPLHCIF", "Path on the CCDB for the GRPLHCIF object"}; + + o2::framework::Configurable phiEntranceCoeff1{"phiEntranceCoeff1", 1.026f, "phiEntrance = phi + phiEntranceCoeff1 * LightSpeedDm2PS * 0.5 * phiEntranceCoeff2 * 1/pT[GeV/c]"}; + o2::framework::Configurable phiEntranceCoeff2{"phiEntranceCoeff2", 85.f, "phiEntrance = phi + phiEntranceCoeff1 * LightSpeedDm2PS * 0.5 * phiEntranceCoeff2 * 1/pT[GeV/c]"}; }; // helper getter - FIXME should be separate @@ -445,10 +453,49 @@ class pidTPCModule template std::vector createNetworkPrediction(TCCDB& ccdb, soa::Join 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 network_prediction; + constexpr std::array 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 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) @@ -486,7 +533,7 @@ class pidTPCModule 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)); @@ -500,49 +547,48 @@ class pidTPCModule } // 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(prediction_size * 9); // For each mass hypotheses + networkPrediction = std::vector(predictionSize * NParticleTypes); // For each mass hypotheses const float nNclNormalization = response->GetNClNormalization(); - float duration_network = 0; + float durationNetwork = 0; - std::vector track_properties(track_prop_size); - uint64_t counter_track_props = 0; - int loop_counter = 0; + std::vector trackProperties(trackPropSize); + uint64_t counterTrackProps = 0; + int loopCounter = 0; // To load the Hadronic rate once for each collision - float hadronicRateBegin = 0.; std::vector 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(); 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 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; @@ -552,74 +598,50 @@ class pidTPCModule 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>(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>(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>(stop_network_total - start_network_total).count() / (size * 9) << "ns ; Total time (eval + overhead): " << std::chrono::duration>(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>(stopNetworkTotal - startNetworkTotal).count() / (size * NParticleTypes) << "ns ; Total time (eval + overhead): " << std::chrono::duration>(stopNetworkTotal - startNetworkTotal).count() / NanoToOne << " s"; - return network_prediction; + return networkPrediction; } //__________________________________________________