Skip to content

Commit

Permalink
[Common,DPG] PID: Add AO2D metadata handling for pass name in TPC CCD…
Browse files Browse the repository at this point in the history 
…B calls (pull/6793), adding NN-version number , using nSigmaTOFdautrack in protons from Lambda decays (#8835)
  • Loading branch information
@amaringarcia
amaringarcia authored Dec 10, 2024
1 parent 70f85cb commit 81c7bcd
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 21 deletions.
74 changes: 56 additions & 18 deletions Common/TableProducer/PID/pidTPC.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,24 +18,32 @@
/// \brief Task to produce PID tables for TPC split for each particle.
/// Only the tables for the mass hypotheses requested are filled, and only for the requested table size ("Full" or "Tiny"). The others are sent empty.
///

#include <utility>
#include <map>
#include <memory>
#include <string>
#include <vector>
// ROOT includes
#include "TFile.h"
#include "TRandom.h"
#include "TSystem.h"

// O2 includes
#include <CCDB/BasicCCDBManager.h>
#include "CCDB/BasicCCDBManager.h"
#include "Framework/AnalysisTask.h"
#include "Framework/runDataProcessing.h"
#include "Framework/ASoAHelpers.h"
#include "ReconstructionDataFormats/Track.h"
#include "CCDB/CcdbApi.h"
#include "Common/DataModel/PIDResponse.h"
#include "Common/Core/PID/TPCPIDResponse.h"
#include "Framework/AnalysisDataModel.h"
#include "Common/DataModel/Multiplicity.h"
#include "Common/DataModel/EventSelection.h"
#include "TableHelper.h"
#include "Tools/ML/model.h"
#include "pidTPCBase.h"
#include "MetadataHelper.h"

using namespace o2;
using namespace o2::framework;
Expand All @@ -45,21 +53,21 @@ using namespace o2::framework::expressions;
using namespace o2::track;
using namespace o2::ml;

MetadataHelper metadataInfo; // Metadata helper

void customize(std::vector<o2::framework::ConfigParamSpec>& workflowOptions)
{
std::vector<ConfigParamSpec> options{{"add-qa", VariantType::Int, 0, {"Legacy. No effect."}}};
std::swap(workflowOptions, options);
}

#include "Framework/runDataProcessing.h"

/// Task to produce the response table
struct tpcPid {
using Trks = soa::Join<aod::Tracks, aod::TracksExtra>;
using Coll = soa::Join<aod::Collisions, aod::PIDMults>;
using Coll = soa::Join<aod::Collisions, aod::PIDMults, aod::EvSels>;

using TrksMC = soa::Join<aod::Tracks, aod::TracksExtra, aod::McTrackLabels>;
using CollMC = soa::Join<aod::Collisions, aod::PIDMults, aod::McCollisionLabels>;
using CollMC = soa::Join<aod::Collisions, aod::PIDMults, aod::McCollisionLabels, aod::EvSels>;

// Tables to produce
Produces<o2::aod::pidTPCFullEl> tablePIDFullEl;
Expand Down Expand Up @@ -90,8 +98,10 @@ struct tpcPid {
OnnxModel network;
o2::ccdb::CcdbApi ccdbApi;
std::map<std::string, std::string> metadata;
std::map<std::string, std::string> nullmetadata;
std::map<std::string, std::string> headers;
std::vector<int> speciesNetworkFlags = std::vector<int>(9);
std::string networkVersion;

// Input parameters
Service<o2::ccdb::BasicCCDBManager> ccdb;
Expand Down Expand Up @@ -187,11 +197,14 @@ struct tpcPid {
speciesNetworkFlags[7] = useNetworkHe;
speciesNetworkFlags[8] = useNetworkAl;

// Initialise metadata object for CCDB calls
// Initialise metadata object for CCDB calls from AO2D metadata
if (recoPass.value == "") {
LOGP(info, "Reco pass not specified; CCDB will take latest available object");
if (metadataInfo.isFullyDefined()) {
metadata["RecoPassName"] = metadataInfo.get("RecoPassName");
LOGP(info, "Automatically setting reco pass for TPC Response to {} from AO2D", metadata["RecoPassName"]);
}
} else {
LOGP(info, "CCDB object will be requested for reconstruction pass {}", recoPass.value);
LOGP(info, "Setting reco pass for TPC response to user-defined name {}", recoPass.value);
metadata["RecoPassName"] = recoPass.value;
}

Expand All @@ -215,17 +228,23 @@ struct tpcPid {
ccdb->setCaching(true);
ccdb->setLocalObjectValidityChecking();
ccdb->setCreatedNotAfter(std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count());
ccdbApi.init(url);
if (time != 0) {
LOGP(info, "Initialising TPC PID response for fixed timestamp {} and reco pass {}:", time, recoPass.value);
ccdb->setTimestamp(time);
response = ccdb->getSpecific<o2::pid::tpc::Response>(path, time, metadata);
headers = ccdbApi.retrieveHeaders(path, metadata, time);
if (!response) {
LOGF(warning, "Unable to find TPC parametrisation for specified pass name - falling back to latest object");
response = ccdb->getForTimeStamp<o2::pid::tpc::Response>(path, time);
headers = ccdbApi.retrieveHeaders(path, metadata, time);
networkVersion = headers["NN-Version"];
if (!response) {
LOGF(fatal, "Unable to find any TPC object corresponding to timestamp {}!", time);
}
}
LOG(info) << "Successfully retrieved TPC PID object from CCDB for timestamp " << time << ", period " << headers["LPMProductionTag"] << ", recoPass " << headers["RecoPassName"];
metadata["RecoPassName"] = headers["RecoPassName"]; // Force pass number for NN request to match retrieved BB
response->PrintAll();
}
}
Expand All @@ -236,19 +255,21 @@ struct tpcPid {
return;
} else {
/// CCDB and auto-fetching
ccdbApi.init(url);

if (!autofetchNetworks) {
if (ccdbTimestamp > 0) {
/// Fetching network for specific timestamp
LOG(info) << "Fetching network for timestamp: " << ccdbTimestamp.value;
bool retrieveSuccess = ccdbApi.retrieveBlob(networkPathCCDB.value, ".", metadata, ccdbTimestamp.value, false, networkPathLocally.value);
headers = ccdbApi.retrieveHeaders(networkPathCCDB.value, metadata, ccdbTimestamp.value);
networkVersion = headers["NN-Version"];
if (retrieveSuccess) {
network.initModel(networkPathLocally.value, enableNetworkOptimizations.value, networkSetNumThreads.value, strtoul(headers["Valid-From"].c_str(), NULL, 0), strtoul(headers["Valid-Until"].c_str(), NULL, 0));
std::vector<float> dummyInput(network.getNumInputNodes(), 1.);
network.evalModel(dummyInput); /// Init the model evaluations
LOGP(info, "Retrieved NN corrections for production tag {}, pass number {}, and NN-Version {}", headers["LPMProductionTag"], headers["RecoPassName"], headers["NN-Version"]);
} else {
LOG(fatal) << "Error encountered while fetching/loading the network from CCDB! Maybe the network doesn't exist yet for this runnumber/timestamp?";
LOG(fatal) << "No valid NN object found matching retrieved Bethe-Bloch parametrisation for pass " << metadata["RecoPassName"] << ". Please ensure that the requested pass has dedicated NN corrections available";
}
} else {
/// Taking the network from local file
Expand All @@ -266,7 +287,7 @@ struct tpcPid {
}
}

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
Partition<Trks> notTPCStandaloneTracks = (aod::track::tpcNClsFindable > static_cast<uint8_t>(0)) && ((aod::track::itsClusterSizes > static_cast<uint32_t>(0)) || (aod::track::trdPattern > static_cast<uint8_t>(0)) || (aod::track::tofExpMom > 0.f && aod::track::tofChi2 > 0.f)); // To count number of tracks for use in NN array
Partition<Trks> tracksWithTPC = (aod::track::tpcNClsFindable > (uint8_t)0);

template <typename C, typename T, typename B>
Expand All @@ -286,26 +307,33 @@ struct tpcPid {
LOGP(info, "Retrieving TPC Response for timestamp {} and recoPass {}:", bc.timestamp(), recoPass.value);
}
response = ccdb->getSpecific<o2::pid::tpc::Response>(ccdbPath.value, bc.timestamp(), metadata);
headers = ccdbApi.retrieveHeaders(ccdbPath.value, metadata, bc.timestamp());
networkVersion = headers["NN-Version"];
if (!response) {
LOGP(warning, "!! Could not find a valid TPC response object for specific pass name {}! Falling back to latest uploaded object.", recoPass.value);
LOGP(warning, "!! Could not find a valid TPC response object for specific pass name {}! Falling back to latest uploaded object.", metadata["RecoPassName"]);
headers = ccdbApi.retrieveHeaders(ccdbPath.value, nullmetadata, bc.timestamp());
response = ccdb->getForTimeStamp<o2::pid::tpc::Response>(ccdbPath.value, bc.timestamp());
if (!response) {
LOGP(fatal, "Could not find ANY TPC response object for the timestamp {}!", bc.timestamp());
}
}
LOG(info) << "Successfully retrieved TPC PID object from CCDB for timestamp " << bc.timestamp() << ", period " << headers["LPMProductionTag"] << ", recoPass " << headers["RecoPassName"];
metadata["RecoPassName"] = headers["RecoPassName"]; // Force pass number for NN request to match retrieved BB
response->PrintAll();
}

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 = ccdbApi.retrieveBlob(networkPathCCDB.value, ".", metadata, bc.timestamp(), false, networkPathLocally.value);
headers = ccdbApi.retrieveHeaders(networkPathCCDB.value, metadata, bc.timestamp());
networkVersion = headers["NN-Version"];
if (retrieveSuccess) {
network.initModel(networkPathLocally.value, enableNetworkOptimizations.value, networkSetNumThreads.value, strtoul(headers["Valid-From"].c_str(), NULL, 0), strtoul(headers["Valid-Until"].c_str(), NULL, 0));
std::vector<float> dummyInput(network.getNumInputNodes(), 1.);
network.evalModel(dummyInput);
LOGP(info, "Retrieved NN corrections for production tag {}, pass number {}, NN-Version number{}", headers["LPMProductionTag"], headers["RecoPassName"], headers["NN-Version"]);
} else {
LOG(fatal) << "Error encountered while fetching/loading the network from CCDB! Maybe the network doesn't exist yet for this runnumber/timestamp?";
LOG(fatal) << "No valid NN object found matching retrieved Bethe-Bloch parametrisation for pass " << metadata["RecoPassName"] << ". Please ensure that the requested pass has dedicated NN corrections available";
}
}
}
Expand Down Expand Up @@ -342,6 +370,9 @@ struct tpcPid {
track_properties[counter_track_props + 3] = o2::track::pid_constants::sMasses[i];
track_properties[counter_track_props + 4] = trk.has_collision() ? collisions.iteratorAt(trk.collisionId()).multTPC() / 11000. : 1.;
track_properties[counter_track_props + 5] = std::sqrt(nNclNormalization / trk.tpcNClsFound());
if (input_dimensions == 7 && networkVersion == "2") {
track_properties[counter_track_props + 6] = trk.has_collision() ? collisions.iteratorAt(trk.collisionId()).ft0cOccupancyInTimeRange() / 60000. : 1.;
}
counter_track_props += input_dimensions;
}

Expand Down Expand Up @@ -483,13 +514,16 @@ struct tpcPid {
LOGP(info, "Retrieving TPC Response for timestamp {} and recoPass {}:", bc.timestamp(), recoPass.value);
}
response = ccdb->getSpecific<o2::pid::tpc::Response>(ccdbPath.value, bc.timestamp(), metadata);
headers = ccdbApi.retrieveHeaders(ccdbPath.value, metadata, bc.timestamp());
if (!response) {
LOGP(warning, "!! Could not find a valid TPC response object for specific pass name {}! Falling back to latest uploaded object.", recoPass.value);
LOGP(warning, "!! Could not find a valid TPC response object for specific pass name {}! Falling back to latest uploaded object.", metadata["RecoPassName"]);
response = ccdb->getForTimeStamp<o2::pid::tpc::Response>(ccdbPath.value, bc.timestamp());
headers = ccdbApi.retrieveHeaders(ccdbPath.value, nullmetadata, bc.timestamp());
if (!response) {
LOGP(fatal, "Could not find ANY TPC response object for the timestamp {}!", bc.timestamp());
}
}
LOG(info) << "Successfully retrieved TPC PID object from CCDB for timestamp " << bc.timestamp() << ", period " << headers["LPMProductionTag"] << ", recoPass " << headers["RecoPassName"];
response->PrintAll();
}

Expand All @@ -515,7 +549,7 @@ struct tpcPid {

PROCESS_SWITCH(tpcPid, processStandard, "Creating PID tables without MC TuneOnData", true);

Partition<TrksMC> mcnotTPCStandaloneTracks = (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
Partition<TrksMC> mcnotTPCStandaloneTracks = (aod::track::tpcNClsFindable > static_cast<uint8_t>(0)) && ((aod::track::itsClusterSizes > static_cast<uint32_t>(0)) || (aod::track::trdPattern > static_cast<uint8_t>(0)) || (aod::track::tofExpMom > 0.f && aod::track::tofChi2 > 0.f)); // To count number of tracks for use in NN array
Partition<TrksMC> mctracksWithTPC = (aod::track::tpcNClsFindable > (uint8_t)0);

void processMcTuneOnData(CollMC const& collisionsMc, TrksMC const& tracksMc, aod::BCsWithTimestamps const& bcs, aod::McParticles const&)
Expand Down Expand Up @@ -573,7 +607,7 @@ struct tpcPid {
}
response = ccdb->getSpecific<o2::pid::tpc::Response>(ccdbPath.value, bc.timestamp(), metadata);
if (!response) {
LOGP(warning, "!! Could not find a valid TPC response object for specific pass name {}! Falling back to latest uploaded object.", recoPass.value);
LOGP(warning, "!! Could not find a valid TPC response object for specific pass name {}! Falling back to latest uploaded object.", metadata["RecoPassName"]);
response = ccdb->getForTimeStamp<o2::pid::tpc::Response>(ccdbPath.value, bc.timestamp());
if (!response) {
LOGP(fatal, "Could not find ANY TPC response object for the timestamp {}!", bc.timestamp());
Expand Down Expand Up @@ -641,4 +675,8 @@ struct tpcPid {
PROCESS_SWITCH(tpcPid, processMcTuneOnData, "Creating PID tables with MC TuneOnData", false);
};

WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) { return WorkflowSpec{adaptAnalysisTask<tpcPid>(cfgc)}; }
WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
{
metadataInfo.initMetadata(cfgc); // Parse AO2D metadata
return WorkflowSpec{adaptAnalysisTask<tpcPid>(cfgc)};
}
10 changes: 7 additions & 3 deletions DPG/Tasks/TPC/tpcSkimsTableCreator.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ struct TreeWriterTpcV0 {
Produces<o2::aod::SkimmedTPCV0Tree> rowTPCTree;

/// Configurables
Configurable<float> nSigmaTOFdautrack{"nSigmaTOFdautrack", 5., "n-sigma TOF cut on the daughter tracks. Set 0 to switch it off."};
Configurable<float> nSigmaTOFdautrack{"nSigmaTOFdautrack", 999., "n-sigma TOF cut on the proton daughter tracks. Set 999 to switch it off."};
Configurable<float> nClNorm{"nClNorm", 152., "Number of cluster normalization. Run 2: 159, Run 3 152"};
Configurable<int> applyEvSel{"applyEvSel", 2, "Flag to apply rapidity cut: 0 -> no event selection, 1 -> Run 2 event selection, 2 -> Run 3 event selection"};
Configurable<int> trackSelection{"trackSelection", 1, "Track selection: 0 -> No Cut, 1 -> kGlobalTrack, 2 -> kGlobalTrackWoPtEta, 3 -> kGlobalTrackWoDCA, 4 -> kQualityTracks, 5 -> kInAcceptanceTracks"};
Expand Down Expand Up @@ -222,7 +222,9 @@ struct TreeWriterTpcV0 {
// Lambda
if (static_cast<bool>(posTrack.pidbit() & (1 << 2)) && static_cast<bool>(negTrack.pidbit() & (1 << 2))) {
if (downsampleTsalisCharged(posTrack.pt(), downsamplingTsalisProtons, sqrtSNN, o2::track::pid_constants::sMasses[o2::track::PID::Proton], maxPt4dwnsmplTsalisProtons)) {
fillSkimmedV0Table(v0, posTrack, collision, posTrack.tpcNSigmaPr(), posTrack.tofNSigmaPr(), posTrack.tpcExpSignalPr(posTrack.tpcSignal()), o2::track::PID::Proton, runnumber, dwnSmplFactor_Pr, hadronicRate);
if (TMath::Abs(posTrack.tofNSigmaPr()) <= nSigmaTOFdautrack) {
fillSkimmedV0Table(v0, posTrack, collision, posTrack.tpcNSigmaPr(), posTrack.tofNSigmaPr(), posTrack.tpcExpSignalPr(posTrack.tpcSignal()), o2::track::PID::Proton, runnumber, dwnSmplFactor_Pr, hadronicRate);
}
}
if (downsampleTsalisCharged(negTrack.pt(), downsamplingTsalisPions, sqrtSNN, o2::track::pid_constants::sMasses[o2::track::PID::Pion], maxPt4dwnsmplTsalisPions)) {
fillSkimmedV0Table(v0, negTrack, collision, negTrack.tpcNSigmaPi(), negTrack.tofNSigmaPi(), negTrack.tpcExpSignalPi(negTrack.tpcSignal()), o2::track::PID::Pion, runnumber, dwnSmplFactor_Pi, hadronicRate);
Expand All @@ -234,7 +236,9 @@ struct TreeWriterTpcV0 {
fillSkimmedV0Table(v0, posTrack, collision, posTrack.tpcNSigmaPi(), posTrack.tofNSigmaPi(), posTrack.tpcExpSignalPi(posTrack.tpcSignal()), o2::track::PID::Pion, runnumber, dwnSmplFactor_Pi, hadronicRate);
}
if (downsampleTsalisCharged(negTrack.pt(), downsamplingTsalisProtons, sqrtSNN, o2::track::pid_constants::sMasses[o2::track::PID::Proton], maxPt4dwnsmplTsalisProtons)) {
fillSkimmedV0Table(v0, negTrack, collision, negTrack.tpcNSigmaPr(), negTrack.tofNSigmaPr(), negTrack.tpcExpSignalPr(negTrack.tpcSignal()), o2::track::PID::Proton, runnumber, dwnSmplFactor_Pr, hadronicRate);
if (TMath::Abs(negTrack.tofNSigmaPr()) <= nSigmaTOFdautrack) {
fillSkimmedV0Table(v0, negTrack, collision, negTrack.tpcNSigmaPr(), negTrack.tofNSigmaPr(), negTrack.tpcExpSignalPr(negTrack.tpcSignal()), o2::track::PID::Proton, runnumber, dwnSmplFactor_Pr, hadronicRate);
}
}
}
}
Expand Down

0 comments on commit 81c7bcd

Please sign in to comment.