diff --git a/Tools/CMakeLists.txt b/Tools/CMakeLists.txt index 2b1de764169..ad870181e8c 100644 --- a/Tools/CMakeLists.txt +++ b/Tools/CMakeLists.txt @@ -12,3 +12,4 @@ add_subdirectory(PIDML) add_subdirectory(ML) add_subdirectory(KFparticle) +add_subdirectory(PIDFeatureExtractor) diff --git a/Tools/PIDFeatureExtractor/CMakeLists.txt b/Tools/PIDFeatureExtractor/CMakeLists.txt new file mode 100644 index 00000000000..b163d0e9293 --- /dev/null +++ b/Tools/PIDFeatureExtractor/CMakeLists.txt @@ -0,0 +1,20 @@ +# Copyright 2019-2020 CERN and copyright holders of ALICE O2. +# See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +# All rights not expressly granted are reserved. +# +# This software is distributed under the terms of the GNU General Public +# License v3 (GPL Version 3), copied verbatim in the file "COPYING". +# +# In applying this license CERN does not waive the privileges and immunities +# granted to it by virtue of its status as an Intergovernmental Organization +# or submit itself to any jurisdiction. + +o2physics_add_dpl_workflow(pid-feature-extractor + SOURCES pidFeatureExtractor.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore + COMPONENT_NAME Analysis) + +o2physics_add_dpl_workflow(pid-onnx-inference + SOURCES pidOnnxInference.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore + COMPONENT_NAME Analysis) diff --git a/Tools/PIDFeatureExtractor/README.md b/Tools/PIDFeatureExtractor/README.md new file mode 100644 index 00000000000..c0773c63fd4 --- /dev/null +++ b/Tools/PIDFeatureExtractor/README.md @@ -0,0 +1,136 @@ +# PID Feature Extractor + ONNX Inference + +This provides particle identification for ALICE Run 3 Pb-Pb analyses using +a trained ML model (a detector-aware attention model conditioned on which +detectors each track actually has hits in - TPC, TOF, TRD, ITS, EMCal, +HMPID, plus event centrality). Two tasks: + +- **`pidFeatureExtractor.cxx`** reads AO2D data and writes out the model's + input features - kinematics, per-detector PID signals, and detector + presence flags - to a ROOT file (and optionally CSV). +- **`pidOnnxInference.cxx`** takes that file, runs the trained ONNX model + over it, and writes back a probability for each particle species + (pion / kaon / proton / electron) per track. + +You run the extractor first, then inference on its output. Both are +regular AOD-subscribing tasks; `PidOnnxInference` does its real work by +reading the extractor's output file directly rather than the AOD data it's +subscribed to (see "Running" below for what that means in practice). + +## PidFeatureExtractor + +An ordinary AOD-subscribing analysis task. It reads track and collision +data and, for each track passing the (optional, off by default) quality +cuts, writes one row containing: + +- kinematics (momentum, eta, phi, DCA) +- per-detector signals for TPC, TOF, TRD, ITS, EMCal, and HMPID, each with + a flag saying whether that detector actually has a hit on this track +- event centrality +- a Bayesian PID posterior, for comparison against the ML model +- for MC only: the true particle ID and whether it's a physical primary + +Mode is a runtime switch - enable `processData` for real data or +`processMc` for MC (reconstructed + truth), not both. + +### PidFeatureExtractor options + +| Option | Default | What it does | +|------------------------|------------------|------------------------------------------------------------------| +| `outputPath` | `pid_features` | Output file base name | +| `exportROOT` | `true` | Write a ROOT file | +| `exportCsv` | `false` | Also write CSV | +| `etaMin` / `etaMax` | `-99` / `99` | Eta cut - wide open by default (no cut) | +| `ptMin` / `ptMax` | `0` / `9999` | pT cut, GeV/c - wide open by default | +| `dcaXYMax` / `dcaZMax` | `9999` / `9999` | DCA cuts, cm - wide open by default | +| `itsMinClusters` | `0` | Minimum ITS clusters - `0` = no cut | +| `tpcMinClusters` | `0` | Minimum TPC clusters - `0` = no cut | +| `computeBayesianPid` | `true` | Compute the comparison Bayesian posterior | +| `bayesianPriors` | flat (`1,1,1,1`) | Per-species priors `[pi, ka, pr, el]` for the Bayesian posterior | + +All the cuts default to "off" - tighten them in your config if you want +quality selection applied here rather than downstream. + +## PidOnnxInference + +Takes the file `PidFeatureExtractor` wrote and runs the trained ONNX model +over it, row by row, in `init()` - not per-collision. The model can be +loaded either from CCDB or from a local file, which is handled by +`o2::analysis::MlResponse` (`Tools/ML/MlResponse.h`). + +This is still a normal AOD-subscribing task, so it needs a valid AO2D +file to run at all, the same as any other task in this repository - but +it doesn't actually use that data; `process()` is intentionally empty. +Point it at any valid AO2D (the same one you ran the extractor against is +the obvious choice) purely to satisfy the pipeline. + +By default it assumes every detector group is present and usable, exactly +as the input data says. If you want to see how the model behaves with a +detector deliberately left out - for testing, or to match a specific +detector configuration - each group can be switched off independently; +turning one off overrides the data for that group, the same way a genuine +detector miss would look. + +### PidOnnxInference options + +| Option | Default | What it does | +|---------------------|-----------------------------|------------------------------------------------------------------------------| +| `inputRootFile` | `pid_features_data.root` | File written by `PidFeatureExtractor` | +| `inputTreeName` | `pid_features` | Tree name inside it | +| `outputPath` | `pid_predictions` | Output file base name | +| `exportCsv` | `false` | Also write CSV | +| `loadModelFromCcdb` | `true` | Load the model from CCDB; set `false` to use a local file instead | +| `ccdbUrl` | `http://alice-ccdb.cern.ch` | | +| `modelPathsCcdb` | *(placeholder)* | CCDB path to your model - set this to a real path before running | +| `timestampCcdb` | `-1` | `-1` = latest | +| `onnxFileNames` | `pid_feature_model.onnx` | Local model file, used when `loadModelFromCcdb` is `false` | +| `useTPC` | `true` | Include TPC. Set `false` to exclude it from inference regardless of the data | +| `useTOF` | `true` | Include TOF | +| `useTRD` | `true` | Include TRD | +| `useITS` | `true` | Include ITS | +| `useEMCal` | `true` | Include EMCal | +| `useHMPID` | `true` | Include HMPID | +| `useCentrality` | `true` | Include event centrality | + +Output columns are `mlProbPi`, `mlProbKa`, `mlProbPr`, `mlProbEl` (one +probability per species) and `mlPredictedClass` (the most likely species, +as an index: `0`=pion, `1`=kaon, `2`=proton, `3`=electron). + +## Running + +Both use the usual `--configuration json://your-config.json` mechanism, +and both are AOD-subscribing tasks - `PidOnnxInference` just doesn't use +the AOD data it's given, it reads `PidFeatureExtractor`'s output file +instead. Run the extractor first: + +```bash +#!/bin/bash + +config_file="my-config.json" + +o2-analysis-timestamp --configuration json://$config_file -b | + o2-analysis-event-selection --configuration json://$config_file -b | + o2-analysis-track-propagation --configuration json://$config_file -b | + o2-analysis-trackselection --configuration json://$config_file -b | + o2-analysis-pid-tpc-base --configuration json://$config_file -b | + o2-analysis-pid-tpc --configuration json://$config_file -b | + o2-analysis-pid-tof-base --configuration json://$config_file -b | + o2-analysis-pid-tof --configuration json://$config_file -b | + o2-analysis-pid-tof-beta --configuration json://$config_file -b | + o2-analysis-multiplicity-table --configuration json://$config_file -b | + o2-analysis-centrality-table --configuration json://$config_file -b | + o2-analysis-pid-feature-extractor --configuration json://$config_file -b +``` + +Then run inference, once the extractor has finished and its output file +exists. Any valid AO2D works as input here, since its content is unused - +reusing the same one is the simplest choice: + +```bash +#!/bin/bash + +config_file="my-config.json" + +o2-analysis-timestamp --configuration json://$config_file -b | + o2-analysis-pid-onnx-inference --configuration json://$config_file -b +``` diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx new file mode 100644 index 00000000000..b7c58ebe9f8 --- /dev/null +++ b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx @@ -0,0 +1,497 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file pidFeatureExtractor.cxx +/// \brief Produce flat, ML-ready PID feature files (ROOT TTree and/or CSV) +/// from ALICE Run 3 Pb-Pb AO2D data, for both MC (reconstructed + +/// truth) and real/raw data. +/// +/// Output is written via manual TFile/TTree/ofstream rather than +/// O2's DECLARE_SOA_TABLE/Produces<> table mechanism. That's a +/// deliberate reversion: two Produces<> tables sharing a column +/// prefix in one struct triggered a reproducible framework-level +/// compile failure (ASoA.h/MetadataTrait constraint-satisfaction +/// errors, and a StructToTuple reflection failure) against this O2 +/// build, independent of table description tag naming. This +/// TFile/TTree approach is the same pattern the original two-file +/// (MC/RAW) version of this task used successfully. +/// +/// \author Robert Forynski + +#include "Common/DataModel/Centrality.h" +#include "Common/DataModel/EventSelection.h" +#include "Common/DataModel/PIDResponseTOF.h" +#include "Common/DataModel/PIDResponseTPC.h" +#include "Common/DataModel/TrackSelectionTables.h" + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace o2; +using namespace o2::framework; +using namespace o2::framework::expressions; + +namespace +{ +constexpr float kNaN = std::numeric_limits::quiet_NaN(); +constexpr int kNumItsLayers = 7; +constexpr int kBitsPerItsLayer = 4; +constexpr uint32_t kItsLayerMask = 0xF; +constexpr int kNumSpecies = 4; // pi, ka, pr, el +constexpr float kEmcalEtaOutOfAcceptance = -900.f; + +/// Detector-presence helpers, local to this project. +template +bool tofMissing(T const& track) +{ + return !track.hasTOF(); +} + +template +bool trdMissing(T const& track) +{ + return !track.hasTRD(); +} + +template +float getTofMass(T const& track) +{ + return tofMissing(track) ? kNaN : track.mass(); +} + +/// itsClusterSizes packs 7 ITS layers into 4 bits each (cluster size per +/// layer, 0 = no hit on that layer). Number of ITS clusters is the count of +/// non-zero nibbles, not the raw column value itself. +template +int getItsNClusters(T const& track) +{ + auto v = static_cast(track.itsClusterSizes()); + int n = 0; + for (int layer = 0; layer < kNumItsLayers; layer++) { + if ((v >> (layer * kBitsPerItsLayer)) & kItsLayerMask) { + n++; + } + } + return n; +} +} // namespace + +/// PidFeatureExtractor: flat PID feature file (ROOT TTree and/or CSV) for +/// ML training/inference. +/// +/// Mode (MC vs. real data) is a runtime PROCESS_SWITCH choice, so one +/// executable serves both use cases. +/// +/// - DPG track cuts (eta/pT/DCA/TPC-cluster/ITS-cluster) are optional and off +/// by default (wide-open ranges) - tighten them in the config if you want +/// them applied here instead of in Python post-processing. +/// - Bayesian PID combination is optional (computeBayesianPid, default +/// true) and configurable priors (bayesianPriors, default flat). Valid +/// whenever TPC is present; TOF is folded in too if also present, but is +/// not required - a TPC-only track still gets a real posterior, not NaN. +/// - Output: ROOT TTree (exportROOT, default true) and/or CSV (exportCsv, +/// default false), written via a single row of member variables bound as +/// TTree branches - not an O2 AOD table. +struct PidFeatureExtractor { + std::unique_ptr outputFile; + std::unique_ptr featureTree; + std::ofstream csvFile; + + // --- output row (bound as TTree branches; also used to build CSV rows) --- + float p = 0, pt = 0, px = 0, py = 0, pz = 0, eta = 0, phi = 0, sign = 0; + int trackType = 0; + float vz = 0, centFT0C = 0, dcaXY = 0, dcaZ = 0; + bool hasTpc = false; + float tpcSignal = 0, tpcNSigmaPi = 0, tpcNSigmaKa = 0, tpcNSigmaPr = 0, tpcNSigmaEl = 0; + int tpcNClsFound = 0; + float tpcChi2NCl = 0; + bool hasTof = false; + float tofMass = 0, beta = 0, tofNSigmaPi = 0, tofNSigmaKa = 0, tofNSigmaPr = 0, tofNSigmaEl = 0; + bool hasTrd = false; + float trdSignal = 0, trdChi2 = 0; + int trdPattern = 0; + int itsClusterSizes = 0; + float itsChi2NCl = 0; + bool hasEmcal = false; + float trackEtaEmcal = 0, trackPhiEmcal = 0; + bool hasHmpid = false; + float hmpidSignal = 0, hmpidQMip = 0; + int hmpidNPhotons = 0, hmpidClusSize = 0; + float hmpidMom = 0; + float bayesProbPi = 0, bayesProbKa = 0, bayesProbPr = 0, bayesProbEl = 0; + int mcPdg = 0; + uint8_t mcIsPhysicalPrimary = 0; + + HistogramRegistry histos{"histos", {}, OutputObjHandlingPolicy::AnalysisObject}; + + Filter trackFilter = requireGlobalTrackInFilter(); + + // --- output configuration -------------------------------------------------- + Configurable outputPath{"outputPath", "pid_features", "Output file base name (no extension)"}; + Configurable exportROOT{"exportROOT", true, "Write a ROOT TTree"}; + Configurable exportCsv{"exportCsv", false, "Also write a CSV, alongside the ROOT output"}; + + // --- DPG cuts: wide-open by default, i.e. effectively disabled ----------- + Configurable etaMin{"etaMin", -99.f, "Minimum track eta (DPG cut; wide-open = disabled)"}; + Configurable etaMax{"etaMax", 99.f, "Maximum track eta (DPG cut; wide-open = disabled)"}; + Configurable ptMin{"ptMin", 0.f, "Minimum track pT, GeV/c (DPG cut; wide-open = disabled)"}; + Configurable ptMax{"ptMax", 9999.f, "Maximum track pT, GeV/c (DPG cut; wide-open = disabled)"}; + Configurable dcaXYMax{"dcaXYMax", 9999.f, "Maximum |DCAxy|, cm (DPG cut; wide-open = disabled)"}; + Configurable dcaZMax{"dcaZMax", 9999.f, "Maximum |DCAz|, cm (DPG cut; wide-open = disabled)"}; + Configurable itsMinClusters{"itsMinClusters", 0, "Minimum number of ITS clusters (DPG cut; 0 = disabled)"}; + Configurable tpcMinClusters{"tpcMinClusters", 0, "Minimum TPC clusters (DPG cut; 0 = disabled)"}; + + // --- Bayesian PID ---------------------------------------------------------- + Configurable computeBayesianPid{"computeBayesianPid", true, "Compute Bayesian PID posteriors (else NaN)"}; + Configurable> bayesianPriors{"bayesianPriors", std::vector{1.f, 1.f, 1.f, 1.f}, "Priors [pi,ka,pr,el]; default flat"}; + + using PidTracks = soa::Filtered>; + + using PidTracksMc = soa::Filtered>; + + using PidCollision = soa::Join::iterator; + + void init(InitContext const&) + { + std::string base = outputPath.value; + + if (exportROOT.value) { + outputFile = std::make_unique((base + ".root").c_str(), "RECREATE"); + featureTree = std::make_unique("pid_features", "PID features"); + + featureTree->Branch("p", &p); + featureTree->Branch("pt", &pt); + featureTree->Branch("px", &px); + featureTree->Branch("py", &py); + featureTree->Branch("pz", &pz); + featureTree->Branch("eta", &eta); + featureTree->Branch("phi", &phi); + featureTree->Branch("sign", &sign); + featureTree->Branch("trackType", &trackType); + featureTree->Branch("vz", &vz); + featureTree->Branch("centFT0C", ¢FT0C); + featureTree->Branch("dcaXY", &dcaXY); + featureTree->Branch("dcaZ", &dcaZ); + featureTree->Branch("hasTPC", &hasTpc); + featureTree->Branch("tpcSignal", &tpcSignal); + featureTree->Branch("tpcNSigmaPi", &tpcNSigmaPi); + featureTree->Branch("tpcNSigmaKa", &tpcNSigmaKa); + featureTree->Branch("tpcNSigmaPr", &tpcNSigmaPr); + featureTree->Branch("tpcNSigmaEl", &tpcNSigmaEl); + featureTree->Branch("tpcNClsFound", &tpcNClsFound); + featureTree->Branch("tpcChi2NCl", &tpcChi2NCl); + featureTree->Branch("hasTOF", &hasTof); + featureTree->Branch("tofMass", &tofMass); + featureTree->Branch("beta", &beta); + featureTree->Branch("tofNSigmaPi", &tofNSigmaPi); + featureTree->Branch("tofNSigmaKa", &tofNSigmaKa); + featureTree->Branch("tofNSigmaPr", &tofNSigmaPr); + featureTree->Branch("tofNSigmaEl", &tofNSigmaEl); + featureTree->Branch("hasTRD", &hasTrd); + featureTree->Branch("trdSignal", &trdSignal); + featureTree->Branch("trdChi2", &trdChi2); + featureTree->Branch("trdPattern", &trdPattern); + featureTree->Branch("itsClusterSizes", &itsClusterSizes); + featureTree->Branch("itsChi2NCl", &itsChi2NCl); + featureTree->Branch("hasEMCal", &hasEmcal); + featureTree->Branch("trackEtaEmcal", &trackEtaEmcal); + featureTree->Branch("trackPhiEmcal", &trackPhiEmcal); + featureTree->Branch("hasHMPID", &hasHmpid); + featureTree->Branch("hmpidSignal", &hmpidSignal); + featureTree->Branch("hmpidQMip", &hmpidQMip); + featureTree->Branch("hmpidNPhotons", &hmpidNPhotons); + featureTree->Branch("hmpidClusSize", &hmpidClusSize); + featureTree->Branch("hmpidMom", &hmpidMom); + featureTree->Branch("bayesProbPi", &bayesProbPi); + featureTree->Branch("bayesProbKa", &bayesProbKa); + featureTree->Branch("bayesProbPr", &bayesProbPr); + featureTree->Branch("bayesProbEl", &bayesProbEl); + if (doprocessMc) { + featureTree->Branch("mcPdg", &mcPdg); + featureTree->Branch("mcIsPhysicalPrimary", &mcIsPhysicalPrimary); + } + } + + if (exportCsv.value) { + csvFile.open(base + (doprocessMc ? "_mc.csv" : "_data.csv")); + csvFile << "p,pt,px,py,pz,eta,phi,sign,trackType," + "vz,centFT0C,dcaXY,dcaZ," + "hasTPC,tpcSignal,tpcNSigmaPi,tpcNSigmaKa,tpcNSigmaPr,tpcNSigmaEl,tpcNClsFound,tpcChi2NCl," + "hasTOF,tofMass,beta,tofNSigmaPi,tofNSigmaKa,tofNSigmaPr,tofNSigmaEl," + "hasTRD,trdSignal,trdChi2,trdPattern," + "itsClusterSizes,itsChi2NCl," + "hasEMCal,trackEtaEmcal,trackPhiEmcal," + "hasHMPID,hmpidSignal,hmpidQMip,hmpidNPhotons,hmpidClusSize,hmpidMom," + "bayesProbPi,bayesProbKa,bayesProbPr,bayesProbEl"; + if (doprocessMc) { + csvFile << ",mcPdg,mcIsPhysicalPrimary"; + } + csvFile << "\n"; + } + + const AxisSpec axisPt{200, 0, 10, "pT"}; + const AxisSpec axisEta{60, -1.5, 1.5, "eta"}; + const AxisSpec axisdEdx{300, 0, 300, "dE/dx"}; + const AxisSpec axisBeta{120, 0, 1.2, "beta"}; + const AxisSpec axisMass{100, -0.2, 2.0, "mass"}; + histos.add("QC/nTracks", "Tracks", kTH1F, {{10000, 0, 100000}}); + histos.add("QC/pt", "pT", kTH1F, {axisPt}); + histos.add("QC/eta", "eta", kTH1F, {axisEta}); + histos.add("QC/tpcDedxVsPt", "dE/dx vs pT", kTH2F, {axisPt, axisdEdx}); + histos.add("QC/tofBetaVsP", "beta vs p", kTH2F, {axisPt, axisBeta}); + histos.add("QC/massVsP", "mass vs p", kTH2F, {axisPt, axisMass}); + } + + /// DPG-style track quality cuts. Wide-open defaults mean this is a no-op + /// unless the config tightens them. + template + bool passesDpgCuts(TTrack const& track) const + { + if (track.pt() < ptMin.value || track.pt() > ptMax.value) + return false; + if (track.eta() < etaMin.value || track.eta() > etaMax.value) + return false; + if (std::abs(track.dcaXY()) > dcaXYMax.value) + return false; + if (std::abs(track.dcaZ()) > dcaZMax.value) + return false; + if (track.tpcNClsFound() < tpcMinClusters.value) + return false; + if (getItsNClusters(track) < itsMinClusters.value) + return false; + return true; + } + + /// Bayesian PID: requires TPC (a TPC-only track still gets a real + /// posterior); folds in TOF too when also present. NaN in all four + /// outputs if TPC is absent or computeBayesianPid is false. + void computeBayesianProbs(bool hasTpcIn, const float nsTPC[4], bool hasTofIn, const float nsTOF[4], float out[4]) const + { + if (!computeBayesianPid.value || !hasTpcIn) { + out[0] = out[1] = out[2] = out[3] = kNaN; + return; + } + auto const& priors = bayesianPriors.value; + float sum = 0.f; + for (int i = 0; i < kNumSpecies; i++) { + float logL = -0.5f * nsTPC[i] * nsTPC[i]; + if (hasTofIn) { + logL += -0.5f * nsTOF[i] * nsTOF[i]; + } + out[i] = std::exp(logL) * priors[i]; + sum += out[i]; + } + for (int i = 0; i < kNumSpecies; i++) { + out[i] = sum > 0.f ? out[i] / sum : 0.25f; + } + } + + /// HMPID is sparse (~0.1% of tracks matched) and linked by track global + /// index rather than joinable 1:1, so it's looked up once per collision + /// instead of per track. + static std::unordered_map buildHmpidMap(aod::HMPIDs const& hmpids) + { + std::unordered_map map; + for (auto h = hmpids.begin(); h != hmpids.end(); ++h) { + map[h.trackId()] = h; + } + return map; + } + + /// Fills the member "output row" for one track. Identical for MC and + /// data - the only thing that differs between the two modes is whether + /// mcPdg/mcIsPhysicalPrimary get set afterwards. + template + void fillRow(TTrack const& track, float vzIn, float centFT0CIn, + std::unordered_map const& hmpidMap) + { + p = track.p(); + pt = track.pt(); + px = track.px(); + py = track.py(); + pz = track.pz(); + eta = track.eta(); + phi = track.phi(); + sign = static_cast(track.sign()); + trackType = track.trackType(); + vz = vzIn; + centFT0C = centFT0CIn; + dcaXY = track.dcaXY(); + dcaZ = track.dcaZ(); + + hasTpc = track.hasTPC(); + tpcSignal = track.tpcSignal(); + tpcNSigmaPi = track.tpcNSigmaPi(); + tpcNSigmaKa = track.tpcNSigmaKa(); + tpcNSigmaPr = track.tpcNSigmaPr(); + tpcNSigmaEl = track.tpcNSigmaEl(); + tpcNClsFound = track.tpcNClsFound(); + tpcChi2NCl = track.tpcChi2NCl(); + + hasTof = !tofMissing(track); + tofMass = getTofMass(track); + beta = track.beta(); + tofNSigmaPi = track.tofNSigmaPi(); + tofNSigmaKa = track.tofNSigmaKa(); + tofNSigmaPr = track.tofNSigmaPr(); + tofNSigmaEl = track.tofNSigmaEl(); + + hasTrd = !trdMissing(track); + trdSignal = track.trdSignal(); + trdChi2 = track.trdChi2(); + trdPattern = track.trdPattern(); + + itsClusterSizes = track.itsClusterSizes(); + itsChi2NCl = track.itsChi2NCl(); + + hasEmcal = track.trackEtaEmcal() > kEmcalEtaOutOfAcceptance; + trackEtaEmcal = track.trackEtaEmcal(); + trackPhiEmcal = track.trackPhiEmcal(); + + hasHmpid = false; + hmpidSignal = kNaN; + hmpidQMip = kNaN; + hmpidNPhotons = 0; + hmpidClusSize = 0; + hmpidMom = kNaN; + if (auto it = hmpidMap.find(track.globalIndex()); it != hmpidMap.end()) { + hasHmpid = true; + hmpidSignal = it->second.hmpidSignal(); + hmpidQMip = it->second.hmpidQMip(); + hmpidNPhotons = it->second.hmpidNPhotons(); + hmpidClusSize = it->second.hmpidClusSize(); + hmpidMom = it->second.hmpidMom(); + } + + float nsTPC[4] = {tpcNSigmaPi, tpcNSigmaKa, tpcNSigmaPr, tpcNSigmaEl}; + float nsTOF[4] = {tofNSigmaPi, tofNSigmaKa, tofNSigmaPr, tofNSigmaEl}; + float bayes[4]; + computeBayesianProbs(hasTpc, nsTPC, hasTof, nsTOF, bayes); + bayesProbPi = bayes[0]; + bayesProbKa = bayes[1]; + bayesProbPr = bayes[2]; + bayesProbEl = bayes[3]; + } + + void fillOutputs() + { + if (exportROOT.value) { + featureTree->Fill(); + } + if (exportCsv.value) { + csvFile << p << ',' << pt << ',' << px << ',' << py << ',' << pz << ',' + << eta << ',' << phi << ',' << sign << ',' << trackType << ',' + << vz << ',' << centFT0C << ',' << dcaXY << ',' << dcaZ << ',' + << hasTpc << ',' << tpcSignal << ',' << tpcNSigmaPi << ',' << tpcNSigmaKa << ',' + << tpcNSigmaPr << ',' << tpcNSigmaEl << ',' << tpcNClsFound << ',' << tpcChi2NCl << ',' + << hasTof << ',' << tofMass << ',' << beta << ',' << tofNSigmaPi << ',' + << tofNSigmaKa << ',' << tofNSigmaPr << ',' << tofNSigmaEl << ',' + << hasTrd << ',' << trdSignal << ',' << trdChi2 << ',' << trdPattern << ',' + << itsClusterSizes << ',' << itsChi2NCl << ',' + << hasEmcal << ',' << trackEtaEmcal << ',' << trackPhiEmcal << ',' + << hasHmpid << ',' << hmpidSignal << ',' << hmpidQMip << ',' + << hmpidNPhotons << ',' << hmpidClusSize << ',' << hmpidMom << ',' + << bayesProbPi << ',' << bayesProbKa << ',' << bayesProbPr << ',' << bayesProbEl; + if (doprocessMc) { + csvFile << ',' << mcPdg << ',' << static_cast(mcIsPhysicalPrimary); + } + csvFile << '\n'; + } + } + + void fillQcHistos() + { + histos.fill(HIST("QC/nTracks"), 1); + histos.fill(HIST("QC/pt"), pt); + histos.fill(HIST("QC/eta"), eta); + if (hasTpc) { + histos.fill(HIST("QC/tpcDedxVsPt"), pt, tpcSignal); + } + if (hasTof) { + histos.fill(HIST("QC/tofBetaVsP"), p, beta); + histos.fill(HIST("QC/massVsP"), p, tofMass); + } + } + + void processData(PidCollision const& collision, PidTracks const& tracks, aod::HMPIDs const& hmpids) + { + auto hmpidMap = buildHmpidMap(hmpids); + for (auto const& track : tracks) { + if (!passesDpgCuts(track)) { + continue; + } + fillRow(track, collision.posZ(), collision.centFT0C(), hmpidMap); + fillOutputs(); + fillQcHistos(); + } + } + PROCESS_SWITCH(PidFeatureExtractor, processData, "Produce PID features for real/raw data (no MC truth)", true); + + void processMc(PidCollision const& collision, PidTracksMc const& tracks, aod::McParticles const&, aod::HMPIDs const& hmpids) + { + auto hmpidMap = buildHmpidMap(hmpids); + for (auto const& track : tracks) { + if (!passesDpgCuts(track)) { + continue; + } + if (!track.has_mcParticle()) { + continue; + } + auto mcParticle = track.mcParticle(); + fillRow(track, collision.posZ(), collision.centFT0C(), hmpidMap); + mcPdg = mcParticle.pdgCode(); + mcIsPhysicalPrimary = static_cast(mcParticle.isPhysicalPrimary()); + fillOutputs(); + fillQcHistos(); + } + } + PROCESS_SWITCH(PidFeatureExtractor, processMc, "Produce PID features for MC (reconstructed + truth)", false); + + void finalize() + { + if (exportROOT.value) { + outputFile->cd(); + featureTree->Write(); + outputFile->Close(); + } + if (exportCsv.value) { + csvFile.close(); + } + } +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +} diff --git a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx new file mode 100644 index 00000000000..86030e3ff82 --- /dev/null +++ b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx @@ -0,0 +1,368 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file pidOnnxInference.cxx +/// \brief Run the FSE PID ONNX model (loaded from CCDB, or a local file for +/// testing) over the ROOT TTree produced by pidFeatureExtractor.cxx, +/// and write per-track class probabilities to a new ROOT file +/// (and/or CSV). +/// +/// A normal Configurable<>-based adaptAnalysisTask, as required by +/// this repository's conventions (workflow topology must be +/// expressed via process function switches / Configurable<>, not +/// hand-built DataProcessorSpec Options{}). All real work still +/// happens once, in init() - it reads the extractor's output file +/// directly via plain TFile/TTree, not an AOD table. process() is +/// intentionally a no-op with a minimal, always-valid AOD argument +/// (aod::Collisions), present only so this registers as a normal +/// analysis task; it does nothing per-collision. Running this task +/// therefore still requires a valid AO2D input to satisfy the +/// pipeline, even though its content is unused. +/// +/// Uses o2::analysis::MlResponse (Tools/ML/MlResponse.h) - +/// O2Physics's generic ONNX/CCDB inference wrapper - for model +/// loading and execution. +/// +/// \author Robert Forynski + +#include "Tools/ML/MlResponse.h" + +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +using namespace o2; +using namespace o2::analysis; +using namespace o2::framework; + +namespace +{ +constexpr int kNumClasses = 4; // pi, ka, pr, el - fixed order throughout, matches the paper's model +constexpr float kNaN = std::numeric_limits::quiet_NaN(); +constexpr int kNumItsLayers = 7; +constexpr int kBitsPerItsLayer = 4; +constexpr uint32_t kItsLayerMask = 0xF; + +/// itsClusterSizes packs 7 ITS layers into 4 bits each; a derived cluster +/// count is a far more sensible model input than the raw packed value. +int getItsNClusters(uint32_t v) +{ + int n = 0; + for (int layer = 0; layer < kNumItsLayers; layer++) { + if ((v >> (layer * kBitsPerItsLayer)) & kItsLayerMask) { + n++; + } + } + return n; +} + +int argmax4(std::vector const& v) +{ + int best = 0; + for (int i = 1; i < kNumClasses; i++) { + if (v[i] > v[best]) { + best = i; + } + } + return best; +} + +/// Per-group enable/disable, independent of what the input tree's hasXXX +/// flags say. Default is everything enabled (true) - the normal case, +/// using each track's real detector coverage as-is. Turning a group off +/// forces its features to the same "absent" sentinel used when the +/// detector genuinely didn't fire, and clears its mask bit - useful for +/// testing how the model behaves with a detector deliberately excluded, +/// independent of the data itself. +struct GroupToggles { + bool useTPC = true; + bool useTOF = true; + bool useTRD = true; + bool useITS = true; + bool useEMCal = true; + bool useHMPID = true; + bool useCentrality = true; +}; + +/// All the real work: load the model, read the whole input tree, run +/// inference row by row, write predictions. Called once from init(). +/// +/// Feature order fed to the model - THIS MUST MATCH YOUR TRAINING SCRIPT'S +/// COLUMN ORDER EXACTLY. Reasonable default (every reconstructed feature +/// except vz/centFT0C/sign/trackType and the Bayesian columns, which are a +/// comparison baseline, not a model input), followed by a 7-length group +/// mask (TPC/TOF/TRD/ITS/EMCal/HMPID/centrality). Each group can be +/// disabled via GroupToggles regardless of what the data says - see there +/// for details. ITS and centrality have no hasXXX flag in the input tree +/// (assumed always-present in the data itself), so their toggle is the +/// only way to exclude them. +void runInference(std::string const& inputRootFile, std::string const& inputTreeName, + std::string const& outputPath, bool exportCsv, + o2::analysis::MlResponse& mlResponse, GroupToggles const& groups) +{ + std::unique_ptr inFile(TFile::Open(inputRootFile.c_str(), "READ")); + if (!inFile || inFile->IsZombie()) { + LOG(fatal) << "Could not open input file " << inputRootFile; + return; + } + auto* tree = dynamic_cast(inFile->Get(inputTreeName.c_str())); + if (!tree) { + LOG(fatal) << "Tree " << inputTreeName << " not found in " << inputRootFile; + return; + } + + // Bind the input branches actually used below - must match + // pidFeatureExtractor.cxx's branch names exactly. + float p = 0, pt = 0, px = 0, py = 0, pz = 0, eta = 0, phi = 0; + float dcaXY = 0, dcaZ = 0; + bool hasTPC = false; + float tpcSignal = 0, tpcNSigmaPi = 0, tpcNSigmaKa = 0, tpcNSigmaPr = 0, tpcNSigmaEl = 0; + int tpcNClsFound = 0; + float tpcChi2NCl = 0; + bool hasTOF = false; + float tofMass = 0, beta = 0, tofNSigmaPi = 0, tofNSigmaKa = 0, tofNSigmaPr = 0, tofNSigmaEl = 0; + bool hasTRD = false; + float trdSignal = 0, trdChi2 = 0; + int trdPattern = 0; + int itsClusterSizes = 0; + float itsChi2NCl = 0; + bool hasEMCal = false; + float trackEtaEmcal = 0, trackPhiEmcal = 0; + bool hasHMPID = false; + float hmpidSignal = 0, hmpidQMip = 0; + int hmpidNPhotons = 0, hmpidClusSize = 0; + float hmpidMom = 0; + + tree->SetBranchAddress("p", &p); + tree->SetBranchAddress("pt", &pt); + tree->SetBranchAddress("px", &px); + tree->SetBranchAddress("py", &py); + tree->SetBranchAddress("pz", &pz); + tree->SetBranchAddress("eta", &eta); + tree->SetBranchAddress("phi", &phi); + tree->SetBranchAddress("dcaXY", &dcaXY); + tree->SetBranchAddress("dcaZ", &dcaZ); + tree->SetBranchAddress("hasTPC", &hasTPC); + tree->SetBranchAddress("tpcSignal", &tpcSignal); + tree->SetBranchAddress("tpcNSigmaPi", &tpcNSigmaPi); + tree->SetBranchAddress("tpcNSigmaKa", &tpcNSigmaKa); + tree->SetBranchAddress("tpcNSigmaPr", &tpcNSigmaPr); + tree->SetBranchAddress("tpcNSigmaEl", &tpcNSigmaEl); + tree->SetBranchAddress("tpcNClsFound", &tpcNClsFound); + tree->SetBranchAddress("tpcChi2NCl", &tpcChi2NCl); + tree->SetBranchAddress("hasTOF", &hasTOF); + tree->SetBranchAddress("tofMass", &tofMass); + tree->SetBranchAddress("beta", &beta); + tree->SetBranchAddress("tofNSigmaPi", &tofNSigmaPi); + tree->SetBranchAddress("tofNSigmaKa", &tofNSigmaKa); + tree->SetBranchAddress("tofNSigmaPr", &tofNSigmaPr); + tree->SetBranchAddress("tofNSigmaEl", &tofNSigmaEl); + tree->SetBranchAddress("hasTRD", &hasTRD); + tree->SetBranchAddress("trdSignal", &trdSignal); + tree->SetBranchAddress("trdChi2", &trdChi2); + tree->SetBranchAddress("trdPattern", &trdPattern); + tree->SetBranchAddress("itsClusterSizes", &itsClusterSizes); + tree->SetBranchAddress("itsChi2NCl", &itsChi2NCl); + tree->SetBranchAddress("hasEMCal", &hasEMCal); + tree->SetBranchAddress("trackEtaEmcal", &trackEtaEmcal); + tree->SetBranchAddress("trackPhiEmcal", &trackPhiEmcal); + tree->SetBranchAddress("hasHMPID", &hasHMPID); + tree->SetBranchAddress("hmpidSignal", &hmpidSignal); + tree->SetBranchAddress("hmpidQMip", &hmpidQMip); + tree->SetBranchAddress("hmpidNPhotons", &hmpidNPhotons); + tree->SetBranchAddress("hmpidClusSize", &hmpidClusSize); + tree->SetBranchAddress("hmpidMom", &hmpidMom); + + std::unique_ptr outFile(TFile::Open((outputPath + ".root").c_str(), "RECREATE")); + TTree outTree("pid_predictions", "PID ML predictions"); + float mlProbPi = 0, mlProbKa = 0, mlProbPr = 0, mlProbEl = 0; + int mlPredictedClass = 0; + outTree.Branch("mlProbPi", &mlProbPi); + outTree.Branch("mlProbKa", &mlProbKa); + outTree.Branch("mlProbPr", &mlProbPr); + outTree.Branch("mlProbEl", &mlProbEl); + outTree.Branch("mlPredictedClass", &mlPredictedClass); + + std::ofstream csv; + if (exportCsv) { + csv.open(outputPath + ".csv"); + csv << "mlProbPi,mlProbKa,mlProbPr,mlProbEl,mlPredictedClass\n"; + } + + std::vector x; + std::vector mlOutput; + Long64_t nEntries = tree->GetEntries(); + for (Long64_t i = 0; i < nEntries; i++) { + tree->GetEntry(i); + + // Effective presence = what the data says AND the group is enabled. + bool effTPC = hasTPC && groups.useTPC; + bool effTOF = hasTOF && groups.useTOF; + bool effTRD = hasTRD && groups.useTRD; + bool effEMCal = hasEMCal && groups.useEMCal; + bool effHMPID = hasHMPID && groups.useHMPID; + + x.clear(); + x.reserve(39 + 7); + x.push_back(p); + x.push_back(pt); + x.push_back(px); + x.push_back(py); + x.push_back(pz); + x.push_back(eta); + x.push_back(phi); + x.push_back(dcaXY); + x.push_back(dcaZ); + x.push_back(static_cast(effTPC)); + x.push_back(effTPC ? tpcSignal : kNaN); + x.push_back(effTPC ? tpcNSigmaPi : kNaN); + x.push_back(effTPC ? tpcNSigmaKa : kNaN); + x.push_back(effTPC ? tpcNSigmaPr : kNaN); + x.push_back(effTPC ? tpcNSigmaEl : kNaN); + x.push_back(effTPC ? static_cast(tpcNClsFound) : 0.f); + x.push_back(effTPC ? tpcChi2NCl : kNaN); + x.push_back(static_cast(effTOF)); + x.push_back(effTOF ? tofMass : kNaN); + x.push_back(effTOF ? beta : kNaN); + x.push_back(effTOF ? tofNSigmaPi : kNaN); + x.push_back(effTOF ? tofNSigmaKa : kNaN); + x.push_back(effTOF ? tofNSigmaPr : kNaN); + x.push_back(effTOF ? tofNSigmaEl : kNaN); + x.push_back(static_cast(effTRD)); + x.push_back(effTRD ? trdSignal : kNaN); + x.push_back(effTRD ? trdChi2 : kNaN); + x.push_back(effTRD ? static_cast(trdPattern) : 0.f); + x.push_back(groups.useITS ? static_cast(getItsNClusters(static_cast(itsClusterSizes))) : 0.f); + x.push_back(groups.useITS ? itsChi2NCl : kNaN); + x.push_back(static_cast(effEMCal)); + x.push_back(effEMCal ? trackEtaEmcal : kNaN); + x.push_back(effEMCal ? trackPhiEmcal : kNaN); + x.push_back(static_cast(effHMPID)); + x.push_back(effHMPID ? hmpidSignal : kNaN); + x.push_back(effHMPID ? hmpidQMip : kNaN); + x.push_back(effHMPID ? static_cast(hmpidNPhotons) : 0.f); + x.push_back(effHMPID ? static_cast(hmpidClusSize) : 0.f); + x.push_back(effHMPID ? hmpidMom : kNaN); + // 7-length group mask + x.push_back(static_cast(effTPC)); + x.push_back(static_cast(effTOF)); + x.push_back(static_cast(effTRD)); + x.push_back(static_cast(groups.useITS)); + x.push_back(static_cast(effEMCal)); + x.push_back(static_cast(effHMPID)); + x.push_back(static_cast(groups.useCentrality)); + + mlResponse.isSelectedMl(x, pt, mlOutput); // return value (selection) unused; mlOutput carries the 4 raw scores + mlProbPi = mlOutput[0]; + mlProbKa = mlOutput[1]; + mlProbPr = mlOutput[2]; + mlProbEl = mlOutput[3]; + mlPredictedClass = argmax4(mlOutput); + outTree.Fill(); + + if (exportCsv) { + csv << mlProbPi << ',' << mlProbKa << ',' << mlProbPr << ',' << mlProbEl << ',' << mlPredictedClass << '\n'; + } + } + + outFile->cd(); + outTree.Write(); + outFile->Close(); + if (exportCsv) { + csv.close(); + } + + LOG(info) << "PidOnnxInference: wrote " << nEntries << " predictions to " << outputPath << ".root"; +} +} // namespace + +/// PidOnnxInference: applies the FSE ONNX model to the features written by +/// pidFeatureExtractor.cxx and writes out per-track class probabilities. +/// +/// All real work happens once, in init() - see runInference() above. +/// process() is intentionally empty; it exists only so this task has a +/// valid AOD-subscribing signature, as required by this repository's +/// conventions. +struct PidOnnxInference { + Configurable inputRootFile{"inputRootFile", "pid_features_data.root", "ROOT file produced by pidFeatureExtractor.cxx"}; + Configurable inputTreeName{"inputTreeName", "pid_features", "Name of the TTree inside inputRootFile"}; + Configurable outputPath{"outputPath", "pid_predictions", "Output file base name (no extension)"}; + Configurable exportCsv{"exportCsv", false, "Also write predictions to CSV alongside the ROOT output"}; + + Configurable loadModelFromCcdb{"loadModelFromCcdb", true, "Load the ONNX model from CCDB (else from onnxFileNames as a local path)"}; + Configurable ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "CCDB URL"}; + Configurable> modelPathsCcdb{"modelPathsCcdb", std::vector{"Users/YOURNAME/PidFeatureExtractor/model"}, "CCDB path to the model"}; + Configurable timestampCcdb{"timestampCcdb", -1, "CCDB query timestamp for the model, -1 = latest"}; + Configurable> onnxFileNames{"onnxFileNames", std::vector{"pid_feature_model.onnx"}, "Local ONNX file path(s), used when loadModelFromCcdb is false"}; + + Configurable> binsPtMl{"binsPtMl", std::vector{-1., 9999.}, "pT bin edges for MlResponse (single bin = model isn't pT-binned)"}; + Configurable nClassesMl{"nClassesMl", static_cast(kNumClasses), "Number of model output classes"}; + + Configurable useTPC{"useTPC", true, "Include TPC in inference. Default true (all detectors present); set false to force TPC excluded regardless of the data"}; + Configurable useTOF{"useTOF", true, "Include TOF in inference"}; + Configurable useTRD{"useTRD", true, "Include TRD in inference"}; + Configurable useITS{"useITS", true, "Include ITS in inference"}; + Configurable useEMCal{"useEMCal", true, "Include EMCal in inference"}; + Configurable useHMPID{"useHMPID", true, "Include HMPID in inference"}; + Configurable useCentrality{"useCentrality", true, "Include centrality in inference"}; + + o2::ccdb::CcdbApi ccdbApi; + o2::analysis::MlResponse mlResponse; + + void init(InitContext&) + { + GroupToggles groups; + groups.useTPC = useTPC.value; + groups.useTOF = useTOF.value; + groups.useTRD = useTRD.value; + groups.useITS = useITS.value; + groups.useEMCal = useEMCal.value; + groups.useHMPID = useHMPID.value; + groups.useCentrality = useCentrality.value; + + // Unused thresholds (CutNot everywhere) - this task always reports + // all four probabilities rather than applying a selection cut, so + // cutsMl/cutDirMl don't need to be user-configurable. + static constexpr double kDefaultCutsMl[1][kNumClasses] = {{0., 0., 0., 0.}}; + LabeledArray cutsMl{kDefaultCutsMl[0], 1, kNumClasses, {"pT bin 0"}, {"prob pi", "prob ka", "prob pr", "prob el"}}; + std::vector cutDirMl{cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot}; + + mlResponse.configure(binsPtMl.value, cutsMl, cutDirMl, static_cast(nClassesMl.value)); + if (loadModelFromCcdb.value) { + ccdbApi.init(ccdbUrl.value); + mlResponse.setModelPathsCCDB(onnxFileNames.value, ccdbApi, modelPathsCcdb.value, timestampCcdb.value); + } else { + mlResponse.setModelPathsLocal(onnxFileNames.value); + } + mlResponse.init(); + + runInference(inputRootFile.value, inputTreeName.value, outputPath.value, exportCsv.value, mlResponse, groups); + } + + /// Intentionally empty - all real work happens once in init(). Present + /// only so this task has a valid, AOD-subscribing process() signature. + void process(aod::Collisions const&) {} +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +}