From b0c94bb618057da35c01d59bea1464dfc69723db Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Thu, 6 Aug 2026 23:18:50 +0100 Subject: [PATCH 01/16] Add PID feature extractor and ONNX inference tasks --- Tools/PIDFeatureExtractor/CMakeLists.txt | 20 + .../PIDFeatureExtractor.cxx | 427 ++++++++++++++++++ .../PIDFeatureExtractor/pidFeatureExtractor.h | 129 ++++++ .../pidFeatureExtractorConfig.json | 150 ++++++ .../PIDFeatureExtractor/pidOnnxInference.cxx | 265 +++++++++++ Tools/PIDFeatureExtractor/run.sh | 77 ++++ 6 files changed, 1068 insertions(+) create mode 100644 Tools/PIDFeatureExtractor/CMakeLists.txt create mode 100644 Tools/PIDFeatureExtractor/PIDFeatureExtractor.cxx create mode 100644 Tools/PIDFeatureExtractor/pidFeatureExtractor.h create mode 100644 Tools/PIDFeatureExtractor/pidFeatureExtractorConfig.json create mode 100644 Tools/PIDFeatureExtractor/pidOnnxInference.cxx create mode 100644 Tools/PIDFeatureExtractor/run.sh diff --git a/Tools/PIDFeatureExtractor/CMakeLists.txt b/Tools/PIDFeatureExtractor/CMakeLists.txt new file mode 100644 index 00000000000..00ec05462cc --- /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) \ No newline at end of file diff --git a/Tools/PIDFeatureExtractor/PIDFeatureExtractor.cxx b/Tools/PIDFeatureExtractor/PIDFeatureExtractor.cxx new file mode 100644 index 00000000000..6b890dd63bf --- /dev/null +++ b/Tools/PIDFeatureExtractor/PIDFeatureExtractor.cxx @@ -0,0 +1,427 @@ +// 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 tables from ALICE Run 3 Pb-Pb +/// AO2D data, for both MC (reconstructed + truth) and real/raw data. +/// DPG track cuts and Bayesian PID are both optional and +/// configuration-driven (off/wide-open by default); CSV export is +/// available alongside the table output for convenience. +/// +/// \author Robert Forynski + +#include "pidFeatureExtractor.h" +// +#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 + +using namespace o2; +using namespace o2::framework; +using namespace o2::framework::expressions; + +namespace +{ +constexpr float kNaN = std::numeric_limits::quiet_NaN(); + +/// 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 < 7; layer++) { + if ((v >> (layer * 4)) & 0xF) { + n++; + } + } + return n; +} + +/// One row of reconstructed (non-MC) features, computed once per track and +/// shared between the table fill and the optional CSV row - keeps the two +/// outputs from ever being able to drift apart. +struct FeatureRow { + float p, pt, px, py, pz, eta, phi, sign; + int trackType; + float vz, centFT0C, dcaXY, dcaZ; + bool hasTpc; + float tpcSignal, tpcNSigmaPi, tpcNSigmaKa, tpcNSigmaPr, tpcNSigmaEl; + int tpcNClsFound; + float tpcChi2NCl; + bool hasTof; + float tofMass, beta, tofNSigmaPi, tofNSigmaKa, tofNSigmaPr, tofNSigmaEl; + bool hasTrd; + float trdSignal, trdChi2; + int trdPattern; + int itsClusterSizes; + float itsChi2NCl; + bool hasEmcal; + float trackEtaEmcal, trackPhiEmcal; + bool hasHmpid; + float hmpidSignal, hmpidQMip; + int hmpidNPhotons, hmpidClusSize; + float hmpidMom; + float bayesProbPi, bayesProbKa, bayesProbPr, bayesProbEl; +}; + +/// CSV header, kept in exactly the same column order as FeatureRow / the +/// reconstructed part of the tables above, so a diff between the ROOT table +/// and the CSV is trivial if the two are ever compared. +constexpr const char* kCsvHeader = + "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"; + +void writeCsvRow(std::ofstream& out, FeatureRow const& r) +{ + out << r.p << ',' << r.pt << ',' << r.px << ',' << r.py << ',' << r.pz << ',' + << r.eta << ',' << r.phi << ',' << r.sign << ',' << r.trackType << ',' + << r.vz << ',' << r.centFT0C << ',' << r.dcaXY << ',' << r.dcaZ << ',' + << r.hasTpc << ',' << r.tpcSignal << ',' << r.tpcNSigmaPi << ',' << r.tpcNSigmaKa << ',' + << r.tpcNSigmaPr << ',' << r.tpcNSigmaEl << ',' << r.tpcNClsFound << ',' << r.tpcChi2NCl << ',' + << r.hasTof << ',' << r.tofMass << ',' << r.beta << ',' << r.tofNSigmaPi << ',' + << r.tofNSigmaKa << ',' << r.tofNSigmaPr << ',' << r.tofNSigmaEl << ',' + << r.hasTrd << ',' << r.trdSignal << ',' << r.trdChi2 << ',' << r.trdPattern << ',' + << r.itsClusterSizes << ',' << r.itsChi2NCl << ',' + << r.hasEmcal << ',' << r.trackEtaEmcal << ',' << r.trackPhiEmcal << ',' + << r.hasHmpid << ',' << r.hmpidSignal << ',' << r.hmpidQMip << ',' + << r.hmpidNPhotons << ',' << r.hmpidClusSize << ',' << r.hmpidMom << ',' + << r.bayesProbPi << ',' << r.bayesProbKa << ',' << r.bayesProbPr << ',' << r.bayesProbEl << '\n'; +} +} // namespace + +/// PidFeatureExtractor: flat PID feature table for ML training/inference. +/// +/// Mode (MC vs. real data) is a runtime PROCESS_SWITCH choice, so one +/// executable and one pair of output tables serve 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. +/// - CSV export is optional (exportCsv, default false) and writes the same +/// reconstructed-feature rows as the ROOT table, alongside it. +/// - No histogramming - add a companion QA task later if needed, rather +/// than folding histograms into this producer. +struct PidFeatureExtractor { + Produces pidFeaturesData; + Produces pidFeaturesMc; + + Filter trackFilter = requireGlobalTrackInFilter(); + + // --- 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", {1.f, 1.f, 1.f, 1.f}, "Priors [pi,ka,pr,el]; default flat"}; + + // --- CSV export -------------------------------------------------------------- + Configurable exportCsv{"exportCsv", false, "Also write reconstructed features to CSV alongside the table"}; + Configurable csvOutputPath{"csvOutputPath", "pid_features", "CSV output file base name (no extension), if exportCsv"}; + + std::ofstream csvFile; + + using PidTracks = soa::Filtered>; + + using PidTracksMc = soa::Filtered>; + + using PidCollision = soa::Join::iterator; + + void init(InitContext const&) + { + if (!exportCsv) { + return; + } + // Only one of processData/processMc is expected to be active; open the + // CSV that matches whichever is (doprocessXxx is generated by + // PROCESS_SWITCH). If both or neither are on, nothing here enforces + // that - see the README note on this. + std::string suffix = doprocessMc ? "_mc.csv" : "_data.csv"; + csvFile.open(csvOutputPath.value + suffix); + csvFile << kCsvHeader << '\n'; + } + + /// 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 || track.pt() > ptMax) + return false; + if (track.eta() < etaMin || track.eta() > etaMax) + 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 hasTpc, float nsTPC[4], bool hasTof, float nsTOF[4], float out[4]) const + { + if (!computeBayesianPid || !hasTpc) { + out[0] = out[1] = out[2] = out[3] = kNaN; + return; + } + auto const& priors = bayesianPriors.value; + float sum = 0.f; + for (int i = 0; i < 4; i++) { + float logL = -0.5f * nsTPC[i] * nsTPC[i]; + if (hasTof) { + logL += -0.5f * nsTOF[i] * nsTOF[i]; + } + out[i] = std::exp(logL) * priors[i]; + sum += out[i]; + } + for (int i = 0; i < 4; 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; + } + + template + struct HmpidRow { + bool has = false; + float signal = kNaN, qMip = kNaN, mom = kNaN; + int nPhotons = 0, clusSize = 0; + + static HmpidRow lookup(TTrack const& track, std::unordered_map const& map) + { + HmpidRow row; + if (auto it = map.find(track.globalIndex()); it != map.end()) { + row.has = true; + row.signal = it->second.hmpidSignal(); + row.qMip = it->second.hmpidQMip(); + row.mom = it->second.hmpidMom(); + row.nPhotons = it->second.hmpidNPhotons(); + row.clusSize = it->second.hmpidClusSize(); + } + return row; + } + }; + + /// Builds the shared reconstructed-feature row for one track. Identical + /// for MC and data - the only thing that differs between the two modes is + /// whether MC-truth columns get appended afterwards. + template + FeatureRow buildFeatureRow(TTrack const& track, float vz, float centFT0C, + std::unordered_map const& hmpidMap) const + { + FeatureRow r{}; + r.p = track.p(); + r.pt = track.pt(); + r.px = track.px(); + r.py = track.py(); + r.pz = track.pz(); + r.eta = track.eta(); + r.phi = track.phi(); + r.sign = static_cast(track.sign()); + r.trackType = track.trackType(); + r.vz = vz; + r.centFT0C = centFT0C; + r.dcaXY = track.dcaXY(); + r.dcaZ = track.dcaZ(); + + r.hasTpc = track.hasTPC(); + r.tpcSignal = track.tpcSignal(); + r.tpcNSigmaPi = track.tpcNSigmaPi(); + r.tpcNSigmaKa = track.tpcNSigmaKa(); + r.tpcNSigmaPr = track.tpcNSigmaPr(); + r.tpcNSigmaEl = track.tpcNSigmaEl(); + r.tpcNClsFound = track.tpcNClsFound(); + r.tpcChi2NCl = track.tpcChi2NCl(); + + r.hasTof = !tofMissing(track); + r.tofMass = getTofMass(track); + r.beta = track.beta(); + r.tofNSigmaPi = track.tofNSigmaPi(); + r.tofNSigmaKa = track.tofNSigmaKa(); + r.tofNSigmaPr = track.tofNSigmaPr(); + r.tofNSigmaEl = track.tofNSigmaEl(); + + r.hasTrd = !trdMissing(track); + r.trdSignal = track.trdSignal(); + r.trdChi2 = track.trdChi2(); + r.trdPattern = track.trdPattern(); + + r.itsClusterSizes = track.itsClusterSizes(); + r.itsChi2NCl = track.itsChi2NCl(); + + r.hasEmcal = track.trackEtaEmcal() > -900.f; + r.trackEtaEmcal = track.trackEtaEmcal(); + r.trackPhiEmcal = track.trackPhiEmcal(); + + auto hm = HmpidRow::lookup(track, hmpidMap); + r.hasHmpid = hm.has; + r.hmpidSignal = hm.signal; + r.hmpidQMip = hm.qMip; + r.hmpidNPhotons = hm.nPhotons; + r.hmpidClusSize = hm.clusSize; + r.hmpidMom = hm.mom; + + float nsTPC[4] = {r.tpcNSigmaPi, r.tpcNSigmaKa, r.tpcNSigmaPr, r.tpcNSigmaEl}; + float nsTOF[4] = {r.tofNSigmaPi, r.tofNSigmaKa, r.tofNSigmaPr, r.tofNSigmaEl}; + float bayes[4]; + computeBayesianProbs(r.hasTpc, nsTPC, r.hasTof, nsTOF, bayes); + r.bayesProbPi = bayes[0]; + r.bayesProbKa = bayes[1]; + r.bayesProbPr = bayes[2]; + r.bayesProbEl = bayes[3]; + + return r; + } + + 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; + } + auto r = buildFeatureRow(track, collision.posZ(), collision.centFT0C(), hmpidMap); + + pidFeaturesData( + r.p, r.pt, r.px, r.py, r.pz, r.eta, r.phi, r.sign, r.trackType, + r.vz, r.centFT0C, r.dcaXY, r.dcaZ, + r.hasTpc, r.tpcSignal, r.tpcNSigmaPi, r.tpcNSigmaKa, r.tpcNSigmaPr, r.tpcNSigmaEl, + r.tpcNClsFound, r.tpcChi2NCl, + r.hasTof, r.tofMass, r.beta, r.tofNSigmaPi, r.tofNSigmaKa, r.tofNSigmaPr, r.tofNSigmaEl, + r.hasTrd, r.trdSignal, r.trdChi2, r.trdPattern, + r.itsClusterSizes, r.itsChi2NCl, + r.hasEmcal, r.trackEtaEmcal, r.trackPhiEmcal, + r.hasHmpid, r.hmpidSignal, r.hmpidQMip, r.hmpidNPhotons, r.hmpidClusSize, r.hmpidMom, + r.bayesProbPi, r.bayesProbKa, r.bayesProbPr, r.bayesProbEl); + + if (exportCsv) { + writeCsvRow(csvFile, r); + } + } + } + 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(); + auto r = buildFeatureRow(track, collision.posZ(), collision.centFT0C(), hmpidMap); + + pidFeaturesMc( + r.p, r.pt, r.px, r.py, r.pz, r.eta, r.phi, r.sign, r.trackType, + r.vz, r.centFT0C, r.dcaXY, r.dcaZ, + r.hasTpc, r.tpcSignal, r.tpcNSigmaPi, r.tpcNSigmaKa, r.tpcNSigmaPr, r.tpcNSigmaEl, + r.tpcNClsFound, r.tpcChi2NCl, + r.hasTof, r.tofMass, r.beta, r.tofNSigmaPi, r.tofNSigmaKa, r.tofNSigmaPr, r.tofNSigmaEl, + r.hasTrd, r.trdSignal, r.trdChi2, r.trdPattern, + r.itsClusterSizes, r.itsChi2NCl, + r.hasEmcal, r.trackEtaEmcal, r.trackPhiEmcal, + r.hasHmpid, r.hmpidSignal, r.hmpidQMip, r.hmpidNPhotons, r.hmpidClusSize, r.hmpidMom, + r.bayesProbPi, r.bayesProbKa, r.bayesProbPr, r.bayesProbEl, + mcParticle.pdgCode(), static_cast(mcParticle.isPhysicalPrimary())); + + if (exportCsv) { + writeCsvRow(csvFile, r); + } + } + } + PROCESS_SWITCH(PidFeatureExtractor, processMc, "Produce PID features for MC (reconstructed + truth)", false); +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +} diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractor.h b/Tools/PIDFeatureExtractor/pidFeatureExtractor.h new file mode 100644 index 00000000000..ff6e0997476 --- /dev/null +++ b/Tools/PIDFeatureExtractor/pidFeatureExtractor.h @@ -0,0 +1,129 @@ +// 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.h +/// \brief Data model for the PID feature extractor: PidFeaturesData / +/// PidFeaturesMc table definitions. Included with a plain quoted +/// filename ("pidFeatureExtractor.h"), which the compiler resolves +/// relative to the including .cxx's own directory first - so this +/// header and every .cxx that uses it must stay in the same folder; +/// no repo-root-relative path to get wrong. +/// +/// \author Robert Forynski + +#ifndef PID_FEATURE_EXTRACTOR_H_ +#define PID_FEATURE_EXTRACTOR_H_ + +#include "Common/DataModel/Centrality.h" +#include "Common/DataModel/PIDResponseTOF.h" +#include "Common/DataModel/PIDResponseTPC.h" +#include "Common/DataModel/TrackSelectionTables.h" + +#include + +#include + +namespace o2::aod +{ +namespace pidfeat +{ +// Kinematics not already stored as static columns upstream (Pt, TrackType, +// DcaXY, DcaZ, TPCSignal, TRDSignal, TRDPattern, TrackEtaEMCAL, +// TrackPhiEMCAL, ITSClusterSizes, ITSChi2NCl, TPCNClsFound are reused +// directly from aod::track:: in the table definitions below). +DECLARE_SOA_COLUMN(P, p, float); //! Track momentum magnitude (GeV/c) +DECLARE_SOA_COLUMN(Px, px, float); //! Track x-momentum (GeV/c) +DECLARE_SOA_COLUMN(Py, py, float); //! Track y-momentum (GeV/c) +DECLARE_SOA_COLUMN(Pz, pz, float); //! Track z-momentum (GeV/c) +DECLARE_SOA_COLUMN(Eta, eta, float); //! Pseudorapidity +DECLARE_SOA_COLUMN(Phi, phi, float); //! Azimuthal angle +DECLARE_SOA_COLUMN(Sign, sign, float); //! Track charge sign + +// Event level, duplicated per track for a self-contained flat table. +DECLARE_SOA_COLUMN(Vz, vz, float); //! Collision vertex z (cm) + +// TOF mass: not exposed as a static column upstream. +DECLARE_SOA_COLUMN(TofMass, tofMass, float); //! TOF-reconstructed mass, NaN if !hasTOF() + +// Detector-presence flags. Kept explicit (rather than inferred solely from +// NaN sentinels) because the downstream FSE/attention model conditions on a +// per-track detector mask. +DECLARE_SOA_COLUMN(HasTPC, hasTPC, uint8_t); //! +DECLARE_SOA_COLUMN(HasTOF, hasTOF, uint8_t); //! +DECLARE_SOA_COLUMN(HasTRD, hasTRD, uint8_t); //! trdPattern() > 0 +DECLARE_SOA_COLUMN(HasEMCal, hasEMCal, uint8_t); //! trackEtaEmcal() in acceptance +DECLARE_SOA_COLUMN(HasHMPID, hasHMPID, uint8_t); //! matched in the sparse HMPID table + +// HMPID: sparse table, matched by hand per collision (see buildHmpidMap() in +// pidFeatureExtractor.cxx). +DECLARE_SOA_COLUMN(HmpidSignal, hmpidSignal, float); //! Cherenkov angle (rad), NaN if !hasHMPID +DECLARE_SOA_COLUMN(HmpidQMip, hmpidQMip, float); //! +DECLARE_SOA_COLUMN(HmpidNPhotons, hmpidNPhotons, int); //! +DECLARE_SOA_COLUMN(HmpidClusSize, hmpidClusSize, int); //! +DECLARE_SOA_COLUMN(HmpidMom, hmpidMom, float); //! + +// Bayesian PID posteriors. NaN when computeBayesianPid is false, or when the +// track doesn't have TPC (see computeBayesianProbs() in +// pidFeatureExtractor.cxx) - never a value that could be mistaken for a real +// posterior. +DECLARE_SOA_COLUMN(BayesProbPi, bayesProbPi, float); //! +DECLARE_SOA_COLUMN(BayesProbKa, bayesProbKa, float); //! +DECLARE_SOA_COLUMN(BayesProbPr, bayesProbPr, float); //! +DECLARE_SOA_COLUMN(BayesProbEl, bayesProbEl, float); //! + +// MC truth. +DECLARE_SOA_COLUMN(IsPhysicalPrimary, isPhysicalPrimary, uint8_t); //! +} // namespace pidfeat + +// Real/raw data: reconstructed features only. +DECLARE_SOA_TABLE(PidFeaturesData, "AOD", "PIDFEATDATA", //! + o2::soa::Index<>, + pidfeat::P, aod::track::Pt, pidfeat::Px, pidfeat::Py, pidfeat::Pz, + pidfeat::Eta, pidfeat::Phi, pidfeat::Sign, aod::track::TrackType, + pidfeat::Vz, aod::cent::CentFT0C, + aod::track::DcaXY, aod::track::DcaZ, + pidfeat::HasTPC, aod::track::TPCSignal, + pidtpc::TPCNSigmaPi, pidtpc::TPCNSigmaKa, pidtpc::TPCNSigmaPr, pidtpc::TPCNSigmaEl, + aod::track::TPCNClsFound, aod::track::TPCChi2NCl, + pidfeat::HasTOF, pidfeat::TofMass, aod::pidtofbeta::Beta, + pidtof::TOFNSigmaPi, pidtof::TOFNSigmaKa, pidtof::TOFNSigmaPr, pidtof::TOFNSigmaEl, + pidfeat::HasTRD, aod::track::TRDSignal, aod::track::TRDChi2, aod::track::TRDPattern, + aod::track::ITSClusterSizes, aod::track::ITSChi2NCl, + pidfeat::HasEMCal, aod::track::TrackEtaEMCAL, aod::track::TrackPhiEMCAL, + pidfeat::HasHMPID, pidfeat::HmpidSignal, pidfeat::HmpidQMip, + pidfeat::HmpidNPhotons, pidfeat::HmpidClusSize, pidfeat::HmpidMom, + pidfeat::BayesProbPi, pidfeat::BayesProbKa, pidfeat::BayesProbPr, pidfeat::BayesProbEl); + +// MC: reconstructed features + truth. Same reconstructed columns as +// PidFeaturesData plus PdgCode/IsPhysicalPrimary - kept as a distinct table +// (rather than optional columns on one table) because O2 tables have a +// fixed schema. +DECLARE_SOA_TABLE(PidFeaturesMc, "AOD", "PIDFEATMC", //! + o2::soa::Index<>, + pidfeat::P, aod::track::Pt, pidfeat::Px, pidfeat::Py, pidfeat::Pz, + pidfeat::Eta, pidfeat::Phi, pidfeat::Sign, aod::track::TrackType, + pidfeat::Vz, aod::cent::CentFT0C, + aod::track::DcaXY, aod::track::DcaZ, + pidfeat::HasTPC, aod::track::TPCSignal, + pidtpc::TPCNSigmaPi, pidtpc::TPCNSigmaKa, pidtpc::TPCNSigmaPr, pidtpc::TPCNSigmaEl, + aod::track::TPCNClsFound, aod::track::TPCChi2NCl, + pidfeat::HasTOF, pidfeat::TofMass, aod::pidtofbeta::Beta, + pidtof::TOFNSigmaPi, pidtof::TOFNSigmaKa, pidtof::TOFNSigmaPr, pidtof::TOFNSigmaEl, + pidfeat::HasTRD, aod::track::TRDSignal, aod::track::TRDChi2, aod::track::TRDPattern, + aod::track::ITSClusterSizes, aod::track::ITSChi2NCl, + pidfeat::HasEMCal, aod::track::TrackEtaEMCAL, aod::track::TrackPhiEMCAL, + pidfeat::HasHMPID, pidfeat::HmpidSignal, pidfeat::HmpidQMip, + pidfeat::HmpidNPhotons, pidfeat::HmpidClusSize, pidfeat::HmpidMom, + pidfeat::BayesProbPi, pidfeat::BayesProbKa, pidfeat::BayesProbPr, pidfeat::BayesProbEl, + aod::mcparticle::PdgCode, pidfeat::IsPhysicalPrimary); +} // namespace o2::aod + +#endif // PID_FEATURE_EXTRACTOR_H_ diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractorConfig.json b/Tools/PIDFeatureExtractor/pidFeatureExtractorConfig.json new file mode 100644 index 00000000000..f221462552a --- /dev/null +++ b/Tools/PIDFeatureExtractor/pidFeatureExtractorConfig.json @@ -0,0 +1,150 @@ +{ + "internal-dpl-clock": "", + + "internal-dpl-aod-reader": { + "time-limit": 0, + "aod-file-private": "AO2D.root", + "orbit-offset-enumeration": 0, + "orbit-multiplier-enumeration": 0, + "start-value-enumeration": 0, + "end-value-enumeration": -1, + "step-value-enumeration": 1 + }, + + "timestamp-task": { + "verbose": 0, + "rct-path": "RCT/Info/RunInformation", + "orbit-reset-path": "CTP/Calib/OrbitReset", + "ccdb-url": "http://alice-ccdb.cern.ch", + "isRun2MC": 0 + }, + + "bc-selection-task": { + "processRun2": 0, + "processRun3": 1 + }, + + "event-selection-task": { + "syst": "PbPb", + "muonSelection": 0, + "customDeltaBC": 0, + "isMC": 0, + "processRun2": 0, + "processRun3": 1 + }, + + "track-propagation": { + "ccdb-url": "http://alice-ccdb.cern.ch", + "grp-path": "GLO/GRP/GRP", + "grp-mag-path": "GLO/Config/GRPMagField", + "mVtxPath": "GLO/Calib/MeanVertex", + "geo-path": "GLO/Config/GeometryAligned", + "useMatLUT": 0, + "processStandard": 1, + "processCovariance": 0, + "processCovarianceMc": 0, + "minPropagationDistance": 83.1 + }, + + "track-selection-task": { + "isRun3": 1 + }, + + "pid-tpc-base": { + "ccdb-url": "http://alice-ccdb.cern.ch", + "parametrization-path": "TPC/Calib/Response", + "parametrization-el-path": "TPC/Calib/ResponseElectron", + "resoPath": "TPC/Calib/PIDResponse", + "ccdb-timestamp": 0, + "useNetworkCorrection": 0, + "autofetch-network": 1, + "enableNetworkOptimization": 1, + "networkPathLocally": "", + "networkPathCCDB": "Analysis/PID/TPC", + "onnxFile": "network.onnx", + "enableNetworkInference": 0 + }, + + "pid-tpc": { + "param-file": "", + "param-sigma": "TPC.PIDResponse.sigma:", + "ccdb-url": "http://alice-ccdb.cern.ch", + "ccdbPath": "TPC/Calib/PIDResponse" + }, + + "pid-tof-base": { + "ccdb-url": "http://alice-ccdb.cern.ch", + "parametrizationPath": "TOF/Calib/Response", + "passName": "", + "timeShiftCCDBPath": "", + "fatalOnPassNotAvailable": 1 + }, + + "pid-tof": { + "param-file": "", + "param-sigma": "TOF.PIDResponse.sigma:", + "ccdb-url": "http://alice-ccdb.cern.ch", + "ccdbPath": "TOF/Calib/Response", + "passName": "", + "timeShiftCCDBPath": "", + "parametrizationPath": "TOF/Calib/Response", + "fatalOnPassNotAvailable": 1 + }, + + "pid-tof-beta": { + "ccdb-url": "http://alice-ccdb.cern.ch" + }, + + "multiplicity-table": { + "doVertexZeq": 1, + "fractionOfEvents": 2, + "processRun2": 0, + "processRun3": 1 + }, + + "centrality-table": { + "ccdb-url": "http://alice-ccdb.cern.ch", + "ccdbPath": "Centrality", + "genName": "", + "processRun2": 0, + "processRun3": 1, + "doNotCrashOnNull": 1, + "processFV0A": 0, + "processFT0M": 0, + "processFT0A": 0, + "processFT0C": 1, + "processFDDM": 0, + "processNTPV": 0, + "processNGlobal": 0, + "processMFT": 0 + }, + + "pid-feature-extractor": { + "processData": 1, + "processMc": 0, + "etaMin": -99.0, + "etaMax": 99.0, + "ptMin": 0.0, + "ptMax": 9999.0, + "dcaxyMax": 9999.0, + "dcazMax": 9999.0, + "tpcMinClusters": 0, + "itsMinClusters": 4, + "computeBayesianPid": true, + "bayesianPriors": [1.0, 1.0, 1.0, 1.0], + "exportCsv": false, + "csvOutputPath": "pid_features" + }, + + "pid-onnx-inference": { + "processData": 1, + "processMc": 0, + "loadModelFromCcdb": true, + "ccdbUrl": "http://alice-ccdb.cern.ch", + "modelPathsCcdb": ["Users/YOURNAME/PidFeatureExtractor/model"], + "timestampCcdb": -1, + "onnxFileNames": ["pid_feature_model.onnx"], + "binsPtMl": [-1.0, 9999.0], + "nClassesMl": 4 + } +} diff --git a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx new file mode 100644 index 00000000000..fa64d223cf2 --- /dev/null +++ b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx @@ -0,0 +1,265 @@ +// 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 tables produced by pidFeatureExtractor.cxx, and +/// write out per-track class probabilities. +/// +/// Uses Tools/ML/MlResponse - O2Physics's generic ONNX/CCDB inference +/// wrapper (not PID-specific, used across several PWG groups) - +/// rather than hand-rolled ONNX/CCDB loading, so this task stays +/// small. It depends only on that shared ML infrastructure and on +/// pidFeatureExtractor.h in this same folder - no other analysis +/// task. +/// +/// \author Robert Forynski + +#include "pidFeatureExtractor.h" +// +#include "Tools/ML/MlResponse.h" + +#include +#include +#include +#include + +#include +#include +#include + +using namespace o2; +using namespace o2::analysis; +using namespace o2::framework; +using namespace o2::framework::expressions; + +// ============================================================================ +// OUTPUT TABLE +// ---------------------------------------------------------------------------- +// Declared directly here since nothing else needs to depend on it. +// ============================================================================ +namespace o2::aod +{ +namespace pidpred +{ +DECLARE_SOA_COLUMN(MlProbPi, mlProbPi, float); //! +DECLARE_SOA_COLUMN(MlProbKa, mlProbKa, float); //! +DECLARE_SOA_COLUMN(MlProbPr, mlProbPr, float); //! +DECLARE_SOA_COLUMN(MlProbEl, mlProbEl, float); //! +DECLARE_SOA_COLUMN(MlPredictedClass, mlPredictedClass, int); //! argmax of the 4 probs above: 0=pi,1=ka,2=pr,3=el +} // namespace pidpred + +DECLARE_SOA_TABLE(PidMlPredictions, "AOD", "PIDMLPRED", //! + o2::soa::Index<>, + pidpred::MlProbPi, pidpred::MlProbKa, pidpred::MlProbPr, pidpred::MlProbEl, + pidpred::MlPredictedClass); +} // namespace o2::aod + +namespace +{ +constexpr int kNumClasses = 4; // pi, ka, pr, el - fixed order throughout, matches the paper's model + +// ---------------------------------------------------------------------------- +// Feature order fed to the model. +// ---------------------------------------------------------------------------- +// THIS MUST MATCH YOUR TRAINING SCRIPT'S COLUMN ORDER EXACTLY - a silent +// mismatch here is the single most likely way this task produces wrong +// predictions without any error. The list below is a reasonable default +// (every reconstructed feature in PidFeaturesData/PidFeaturesMc except +// vz/centFT0C/sign/trackType and the Bayesian columns, which are a +// comparison baseline, not a model input) - it is NOT verified against your +// actual training code. Reorder, add, or drop entries to match exactly +// before trusting the output. +// +// If your ONNX export takes features and mask as two separate input +// tensors rather than one concatenated vector, split this function +// accordingly. + +/// 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. +/// Small and duplicated here rather than shared with pidFeatureExtractor.cxx +/// so this task stays self-contained. +template +int getItsNClusters(TRow const& row) +{ + auto v = static_cast(row.itsClusterSizes()); + int n = 0; + for (int layer = 0; layer < 7; layer++) { + if ((v >> (layer * 4)) & 0xF) { + n++; + } + } + return n; +} + +template +std::vector buildModelInput(TRow const& row) +{ + std::vector x; + x.reserve(38 + 7); + + // Kinematics + x.push_back(row.p()); + x.push_back(row.pt()); + x.push_back(row.px()); + x.push_back(row.py()); + x.push_back(row.pz()); + x.push_back(row.eta()); + x.push_back(row.phi()); + // Impact parameters + x.push_back(row.dcaXY()); + x.push_back(row.dcaZ()); + // TPC + x.push_back(static_cast(row.hasTPC())); + x.push_back(row.tpcSignal()); + x.push_back(row.tpcNSigmaPi()); + x.push_back(row.tpcNSigmaKa()); + x.push_back(row.tpcNSigmaPr()); + x.push_back(row.tpcNSigmaEl()); + x.push_back(static_cast(row.tpcNClsFound())); + x.push_back(row.tpcChi2NCl()); + // TOF + x.push_back(static_cast(row.hasTOF())); + x.push_back(row.tofMass()); + x.push_back(row.beta()); + x.push_back(row.tofNSigmaPi()); + x.push_back(row.tofNSigmaKa()); + x.push_back(row.tofNSigmaPr()); + x.push_back(row.tofNSigmaEl()); + // TRD + x.push_back(static_cast(row.hasTRD())); + x.push_back(row.trdSignal()); + x.push_back(row.trdChi2()); + x.push_back(static_cast(row.trdPattern())); + // ITS + x.push_back(static_cast(getItsNClusters(row))); + x.push_back(row.itsChi2NCl()); + // EMCal + x.push_back(static_cast(row.hasEMCal())); + x.push_back(row.trackEtaEmcal()); + x.push_back(row.trackPhiEmcal()); + // HMPID + x.push_back(static_cast(row.hasHMPID())); + x.push_back(row.hmpidSignal()); + x.push_back(row.hmpidQMip()); + x.push_back(static_cast(row.hmpidNPhotons())); + x.push_back(static_cast(row.hmpidClusSize())); + x.push_back(row.hmpidMom()); + + // 7-length group mask: TPC, TOF, TRD, ITS, EMCal, HMPID, centrality. + // ITS is assumed always present (global track requires it); centrality is + // assumed always present. Both are real assumptions, not derived facts - + // adjust if your training data ever has either group absent. + x.push_back(static_cast(row.hasTPC())); + x.push_back(static_cast(row.hasTOF())); + x.push_back(static_cast(row.hasTRD())); + x.push_back(1.f); // ITS + x.push_back(static_cast(row.hasEMCal())); + x.push_back(static_cast(row.hasHMPID())); + x.push_back(1.f); // centrality + + return x; +} + +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; +} +} // namespace + +/// PidOnnxInference: applies the FSE ONNX model to features produced by +/// pidFeatureExtractor.cxx and writes out per-track class probabilities. +/// +/// Model loading (local file or CCDB) and ONNX execution are entirely +/// handled by o2::analysis::MlResponse - this task only builds the input +/// feature vector and reads back the output. A single global pT bin is used +/// by default since the paper's model isn't pT-binned; MlResponse's usual +/// per-class cut mechanism is disabled (cutDirMl = CutNot for every class, +/// confirmed value 2 in the real o2::cuts_ml enum) - this task always +/// reports all four probabilities rather than applying a pass/fail cut. +/// +/// Two things MlResponse enforces that are worth knowing before debugging a +/// failure here: getModelOutput() calls LOG(fatal) if the input vector's +/// length doesn't match the ONNX model's declared input node count (unless +/// that node is a dynamic axis), and separately if pt lands outside +/// binsPtMl's range entirely (see the comment on binsPtMl below). +struct PidOnnxInference { + Produces pidMlPredictions; + + // --- model location ---------------------------------------------------------- + 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"}; + + // --- MlResponse plumbing: a single pT bin, no selection cut applied -------- + // Lower edge is -1 (not 0): MlResponse::findBin() rejects value < front() + // as out-of-range (fatal in getModelOutput), and track.pt() could in + // principle be exactly 0 - keeping the edge below any physical pT avoids + // that boundary case entirely. + Configurable> binsPtMl{"binsPtMl", std::vector{-1., 9999.}, "pT bin edges for MlResponse (single bin = model isn't pT-binned)"}; + Configurable> cutDirMl{"cutDirMl", std::vector{cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot}, "Per-class cut direction; CutNot = always accept, this task doesn't select"}; + + static constexpr double kDefaultCutsMl[1][kNumClasses] = {{0., 0., 0., 0.}}; + Configurable> cutsMl{"cutsMl", {kDefaultCutsMl[0], 1, kNumClasses, {"pT bin 0"}, {"prob pi", "prob ka", "prob pr", "prob el"}}, "Unused thresholds (CutNot everywhere) - required by MlResponse's interface"}; + Configurable nClassesMl{"nClassesMl", static_cast(kNumClasses), "Number of model output classes"}; + + o2::ccdb::CcdbApi ccdbApi; + o2::analysis::MlResponse mlResponse; + std::vector mlOutput; + + void init(InitContext const&) + { + mlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); + if (loadModelFromCcdb) { + ccdbApi.init(ccdbUrl.value); + mlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCcdb.value, timestampCcdb.value); + } else { + mlResponse.setModelPathsLocal(onnxFileNames); + } + mlResponse.init(); + } + + template + void runInference(TTable const& rows) + { + for (auto const& row : rows) { + auto x = buildModelInput(row); + mlResponse.isSelectedMl(x, row.pt(), mlOutput); // return value (selection) unused; mlOutput carries the 4 raw scores + pidMlPredictions(mlOutput[0], mlOutput[1], mlOutput[2], mlOutput[3], argmax4(mlOutput)); + mlOutput.clear(); + } + } + + void processData(aod::PidFeaturesData const& rows) + { + runInference(rows); + } + PROCESS_SWITCH(PidOnnxInference, processData, "Run inference on PidFeaturesData", true); + + void processMc(aod::PidFeaturesMc const& rows) + { + runInference(rows); + } + PROCESS_SWITCH(PidOnnxInference, processMc, "Run inference on PidFeaturesMc", false); +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +} diff --git a/Tools/PIDFeatureExtractor/run.sh b/Tools/PIDFeatureExtractor/run.sh new file mode 100644 index 00000000000..41bd0e16243 --- /dev/null +++ b/Tools/PIDFeatureExtractor/run.sh @@ -0,0 +1,77 @@ +#!/bin/bash + +# PID Feature Extractor Workflow (simplified, table-based) +# Detectors: TPC, TOF, TRD, ITS, EMCal, HMPID + centrality (FT0C) - the 7 +# detector groups / 34-feature contract used by the FSE PID model. +# Event level: Centrality FT0C (Pb-Pb Run 3) +# +# Output is two AOD-joinable tables (PidFeaturesData / PidFeaturesMc) +# written via the framework's own Produces<> mechanism; set the usual +# --aod-writer-json if you want to control where the output +# AnalysisResults-style file goes. CSV export alongside the table is +# available via exportCsv (off by default) - see pidFeatureExtractorConfig.json. +# +# DPG cuts (eta/pT/DCA/TPC-cluster) and Bayesian PID (TPC alone is enough; +# TOF folded in if present; priors configurable) are both optional, set in +# the same config block - see README.md for details. +# +# Mode is a JSON choice: pidFeatureExtractorConfig.json -> "pid-feature-extractor" +# "processData": 1, "processMc": 0 -> real/raw data (no MC truth) +# "processData": 0, "processMc": 1 -> MC (reconstructed + truth) +# This task does not itself abort if you leave both/neither on - the +# framework will just run whichever you've enabled or do nothing productive +# if neither is. Double check your config. + +CONFIG="$(pwd)/pidFeatureExtractorConfig.json" +OPTION="-b --configuration json://${CONFIG}" + +# CRITICAL: Add shared memory flag (~half your available system RAM) +SHM_SIZE="--shm-segment-size 4000000000" + +EXTRACTOR=~/alice/sw/BUILD/O2Physics-latest/O2Physics/stage/bin/o2-analysis-pid-feature-extractor +INFERENCE=~/alice/sw/BUILD/O2Physics-latest/O2Physics/stage/bin/o2-analysis-pid-onnx-inference + +echo "Starting O2Physics PID Feature Extraction + Inference Workflow..." +echo "Using configuration: ${CONFIG}" +echo "Shared memory segment size: ${SHM_SIZE}" +echo "" + +# Pipeline: +# timestamp → event selection → track propagation → track selection +# (needed for requireGlobalTrackInFilter()) +# → TPC PID → TOF PID → TOF beta +# → multiplicity → centrality (needed for CentFT0Cs, both MC and data) +# → feature extractor → ONNX inference +# +# TRD, ITS, EMCal: already in TracksExtra — no extra task needed +# HMPID: already in AO2D O2hmpid_001 — no extra task needed +# +# The inference task consumes PidFeaturesData/PidFeaturesMc directly (DPL +# wires the table dependency automatically since both tasks run in one +# workflow here) - drop the last pipe stage if you only want the features +# table and don't want to run inference. + +o2-analysis-timestamp ${OPTION} ${SHM_SIZE} | \ +o2-analysis-event-selection ${OPTION} ${SHM_SIZE} | \ +o2-analysis-track-propagation ${OPTION} ${SHM_SIZE} | \ +o2-analysis-trackselection ${OPTION} ${SHM_SIZE} | \ +o2-analysis-pid-tpc-base ${OPTION} ${SHM_SIZE} | \ +o2-analysis-pid-tpc ${OPTION} ${SHM_SIZE} | \ +o2-analysis-pid-tof-base ${OPTION} ${SHM_SIZE} | \ +o2-analysis-pid-tof ${OPTION} ${SHM_SIZE} | \ +o2-analysis-pid-tof-beta ${OPTION} ${SHM_SIZE} | \ +o2-analysis-multiplicity-table ${OPTION} ${SHM_SIZE} | \ +o2-analysis-centrality-table ${OPTION} ${SHM_SIZE} | \ +${EXTRACTOR} ${OPTION} ${SHM_SIZE} | \ +${INFERENCE} ${OPTION} ${SHM_SIZE} + +EXIT_CODE=$? + +echo "" +if [ $EXIT_CODE -eq 0 ]; then + echo "✓ Workflow completed successfully!" +else + echo "✗ Workflow failed with exit code: $EXIT_CODE" +fi + +exit $EXIT_CODE From d06d1844b75cd6aa03b3412a149aa7b84f3b3765 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 16:21:17 +0100 Subject: [PATCH 02/16] Fix case-only filename mismatch (PIDFeatureExtractor.cxx -> pidFeatureExtractor.cxx) --- .../{PIDFeatureExtractor.cxx => pidFeatureExtractor.cxx} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename Tools/PIDFeatureExtractor/{PIDFeatureExtractor.cxx => pidFeatureExtractor.cxx} (100%) diff --git a/Tools/PIDFeatureExtractor/PIDFeatureExtractor.cxx b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx similarity index 100% rename from Tools/PIDFeatureExtractor/PIDFeatureExtractor.cxx rename to Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx From 735c12397b8f36719ea86d7bb65631a889ac3c79 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 16:31:55 +0100 Subject: [PATCH 03/16] Register PIDFeatureExtractor subdirectory in Tools/CMakeLists.txt --- Tools/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) 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) From f4f8c99e21d0fc77e5c2c27ede17f17542ff53ed Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 16:43:19 +0100 Subject: [PATCH 04/16] Fix Configurable comparison operators and ofstream reflection issue --- Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx index 6b890dd63bf..06139a1d653 100644 --- a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx +++ b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -181,7 +182,7 @@ struct PidFeatureExtractor { Configurable exportCsv{"exportCsv", false, "Also write reconstructed features to CSV alongside the table"}; Configurable csvOutputPath{"csvOutputPath", "pid_features", "CSV output file base name (no extension), if exportCsv"}; - std::ofstream csvFile; + std::unique_ptr csvFile; using PidTracks = soa::Filtered(csvOutputPath.value + suffix); + *csvFile << kCsvHeader << '\n'; } /// DPG-style track quality cuts. Wide-open defaults mean this is a no-op @@ -216,9 +217,9 @@ struct PidFeatureExtractor { template bool passesDpgCuts(TTrack const& track) const { - if (track.pt() < ptMin || track.pt() > ptMax) + if (track.pt() < ptMin.value || track.pt() > ptMax.value) return false; - if (track.eta() < etaMin || track.eta() > etaMax) + if (track.eta() < etaMin.value || track.eta() > etaMax.value) return false; if (std::abs(track.dcaXY()) > dcaXYMax.value) return false; @@ -381,7 +382,7 @@ struct PidFeatureExtractor { r.bayesProbPi, r.bayesProbKa, r.bayesProbPr, r.bayesProbEl); if (exportCsv) { - writeCsvRow(csvFile, r); + writeCsvRow(*csvFile, r); } } } @@ -414,7 +415,7 @@ struct PidFeatureExtractor { mcParticle.pdgCode(), static_cast(mcParticle.isPhysicalPrimary())); if (exportCsv) { - writeCsvRow(csvFile, r); + writeCsvRow(*csvFile, r); } } } From 62e78d10d06c8cab39daa09f83253db938dda015 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 16:48:49 +0100 Subject: [PATCH 05/16] Fix remaining .value issues on Configurable comparisons/negation --- Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx index 06139a1d653..2e543cb1b32 100644 --- a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx +++ b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx @@ -200,7 +200,7 @@ struct PidFeatureExtractor { void init(InitContext const&) { - if (!exportCsv) { + if (!exportCsv.value) { return; } // Only one of processData/processMc is expected to be active; open the @@ -237,7 +237,7 @@ struct PidFeatureExtractor { /// outputs if TPC is absent or computeBayesianPid is false. void computeBayesianProbs(bool hasTpc, float nsTPC[4], bool hasTof, float nsTOF[4], float out[4]) const { - if (!computeBayesianPid || !hasTpc) { + if (!computeBayesianPid.value || !hasTpc) { out[0] = out[1] = out[2] = out[3] = kNaN; return; } @@ -381,7 +381,7 @@ struct PidFeatureExtractor { r.hasHmpid, r.hmpidSignal, r.hmpidQMip, r.hmpidNPhotons, r.hmpidClusSize, r.hmpidMom, r.bayesProbPi, r.bayesProbKa, r.bayesProbPr, r.bayesProbEl); - if (exportCsv) { + if (exportCsv.value) { writeCsvRow(*csvFile, r); } } @@ -414,7 +414,7 @@ struct PidFeatureExtractor { r.bayesProbPi, r.bayesProbKa, r.bayesProbPr, r.bayesProbEl, mcParticle.pdgCode(), static_cast(mcParticle.isPhysicalPrimary())); - if (exportCsv) { + if (exportCsv.value) { writeCsvRow(*csvFile, r); } } From 4a5cc974d3df867e658da954f792bba988adbc53 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 17:05:34 +0100 Subject: [PATCH 06/16] Use maximally distinct table description tags to rule out a metadata collision --- Tools/PIDFeatureExtractor/pidFeatureExtractor.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractor.h b/Tools/PIDFeatureExtractor/pidFeatureExtractor.h index ff6e0997476..57ea3aae774 100644 --- a/Tools/PIDFeatureExtractor/pidFeatureExtractor.h +++ b/Tools/PIDFeatureExtractor/pidFeatureExtractor.h @@ -84,7 +84,7 @@ DECLARE_SOA_COLUMN(IsPhysicalPrimary, isPhysicalPrimary, uint8_t); //! } // namespace pidfeat // Real/raw data: reconstructed features only. -DECLARE_SOA_TABLE(PidFeaturesData, "AOD", "PIDFEATDATA", //! +DECLARE_SOA_TABLE(PidFeaturesData, "AOD", "PIDFEATD", //! o2::soa::Index<>, pidfeat::P, aod::track::Pt, pidfeat::Px, pidfeat::Py, pidfeat::Pz, pidfeat::Eta, pidfeat::Phi, pidfeat::Sign, aod::track::TrackType, @@ -106,7 +106,7 @@ DECLARE_SOA_TABLE(PidFeaturesData, "AOD", "PIDFEATDATA", //! // PidFeaturesData plus PdgCode/IsPhysicalPrimary - kept as a distinct table // (rather than optional columns on one table) because O2 tables have a // fixed schema. -DECLARE_SOA_TABLE(PidFeaturesMc, "AOD", "PIDFEATMC", //! +DECLARE_SOA_TABLE(PidFeaturesMc, "AOD", "MCPIDFEA", //! o2::soa::Index<>, pidfeat::P, aod::track::Pt, pidfeat::Px, pidfeat::Py, pidfeat::Pz, pidfeat::Eta, pidfeat::Phi, pidfeat::Sign, aod::track::TrackType, From 4e972583654c95f42d9fb997df503fe81f02ee61 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 17:17:59 +0100 Subject: [PATCH 07/16] Revert to manual TFile/TTree/CSV output, avoiding the Produces<>/DECLARE_SOA_TABLE framework issue --- .../pidFeatureExtractor.cxx | 492 ++++++++++-------- 1 file changed, 278 insertions(+), 214 deletions(-) diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx index 2e543cb1b32..78b61e3c681 100644 --- a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx +++ b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx @@ -10,16 +10,22 @@ // or submit itself to any jurisdiction. /// \file pidFeatureExtractor.cxx -/// \brief Produce flat, ML-ready PID feature tables from ALICE Run 3 Pb-Pb -/// AO2D data, for both MC (reconstructed + truth) and real/raw data. -/// DPG track cuts and Bayesian PID are both optional and -/// configuration-driven (off/wide-open by default); CSV export is -/// available alongside the table output for convenience. +/// \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 "pidFeatureExtractor.h" -// #include "Common/DataModel/Centrality.h" #include "Common/DataModel/EventSelection.h" #include "Common/DataModel/PIDResponseTOF.h" @@ -29,8 +35,12 @@ #include #include #include +#include #include +#include +#include + #include #include #include @@ -82,88 +92,63 @@ int getItsNClusters(T const& track) } return n; } - -/// One row of reconstructed (non-MC) features, computed once per track and -/// shared between the table fill and the optional CSV row - keeps the two -/// outputs from ever being able to drift apart. -struct FeatureRow { - float p, pt, px, py, pz, eta, phi, sign; - int trackType; - float vz, centFT0C, dcaXY, dcaZ; - bool hasTpc; - float tpcSignal, tpcNSigmaPi, tpcNSigmaKa, tpcNSigmaPr, tpcNSigmaEl; - int tpcNClsFound; - float tpcChi2NCl; - bool hasTof; - float tofMass, beta, tofNSigmaPi, tofNSigmaKa, tofNSigmaPr, tofNSigmaEl; - bool hasTrd; - float trdSignal, trdChi2; - int trdPattern; - int itsClusterSizes; - float itsChi2NCl; - bool hasEmcal; - float trackEtaEmcal, trackPhiEmcal; - bool hasHmpid; - float hmpidSignal, hmpidQMip; - int hmpidNPhotons, hmpidClusSize; - float hmpidMom; - float bayesProbPi, bayesProbKa, bayesProbPr, bayesProbEl; -}; - -/// CSV header, kept in exactly the same column order as FeatureRow / the -/// reconstructed part of the tables above, so a diff between the ROOT table -/// and the CSV is trivial if the two are ever compared. -constexpr const char* kCsvHeader = - "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"; - -void writeCsvRow(std::ofstream& out, FeatureRow const& r) -{ - out << r.p << ',' << r.pt << ',' << r.px << ',' << r.py << ',' << r.pz << ',' - << r.eta << ',' << r.phi << ',' << r.sign << ',' << r.trackType << ',' - << r.vz << ',' << r.centFT0C << ',' << r.dcaXY << ',' << r.dcaZ << ',' - << r.hasTpc << ',' << r.tpcSignal << ',' << r.tpcNSigmaPi << ',' << r.tpcNSigmaKa << ',' - << r.tpcNSigmaPr << ',' << r.tpcNSigmaEl << ',' << r.tpcNClsFound << ',' << r.tpcChi2NCl << ',' - << r.hasTof << ',' << r.tofMass << ',' << r.beta << ',' << r.tofNSigmaPi << ',' - << r.tofNSigmaKa << ',' << r.tofNSigmaPr << ',' << r.tofNSigmaEl << ',' - << r.hasTrd << ',' << r.trdSignal << ',' << r.trdChi2 << ',' << r.trdPattern << ',' - << r.itsClusterSizes << ',' << r.itsChi2NCl << ',' - << r.hasEmcal << ',' << r.trackEtaEmcal << ',' << r.trackPhiEmcal << ',' - << r.hasHmpid << ',' << r.hmpidSignal << ',' << r.hmpidQMip << ',' - << r.hmpidNPhotons << ',' << r.hmpidClusSize << ',' << r.hmpidMom << ',' - << r.bayesProbPi << ',' << r.bayesProbKa << ',' << r.bayesProbPr << ',' << r.bayesProbEl << '\n'; -} } // namespace -/// PidFeatureExtractor: flat PID feature table for ML training/inference. +/// 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 and one pair of output tables serve both use cases. +/// 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 +/// - 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. -/// - CSV export is optional (exportCsv, default false) and writes the same -/// reconstructed-feature rows as the ROOT table, alongside it. -/// - No histogramming - add a companion QA task later if needed, rather -/// than folding histograms into this producer. +/// - 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 { - Produces pidFeaturesData; - Produces pidFeaturesMc; + 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)"}; @@ -176,13 +161,7 @@ struct PidFeatureExtractor { // --- Bayesian PID ---------------------------------------------------------- Configurable computeBayesianPid{"computeBayesianPid", true, "Compute Bayesian PID posteriors (else NaN)"}; - Configurable> bayesianPriors{"bayesianPriors", {1.f, 1.f, 1.f, 1.f}, "Priors [pi,ka,pr,el]; default flat"}; - - // --- CSV export -------------------------------------------------------------- - Configurable exportCsv{"exportCsv", false, "Also write reconstructed features to CSV alongside the table"}; - Configurable csvOutputPath{"csvOutputPath", "pid_features", "CSV output file base name (no extension), if exportCsv"}; - - std::unique_ptr csvFile; + Configurable> bayesianPriors{"bayesianPriors", std::vector{1.f, 1.f, 1.f, 1.f}, "Priors [pi,ka,pr,el]; default flat"}; using PidTracks = soa::Filtered((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); + } } - // Only one of processData/processMc is expected to be active; open the - // CSV that matches whichever is (doprocessXxx is generated by - // PROCESS_SWITCH). If both or neither are on, nothing here enforces - // that - see the README note on this. - std::string suffix = doprocessMc ? "_mc.csv" : "_data.csv"; - csvFile = std::make_unique(csvOutputPath.value + suffix); - *csvFile << kCsvHeader << '\n'; + + 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 @@ -235,9 +291,9 @@ struct PidFeatureExtractor { /// 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 hasTpc, float nsTPC[4], bool hasTof, float nsTOF[4], float out[4]) const + void computeBayesianProbs(bool hasTpcIn, float nsTPC[4], bool hasTofIn, float nsTOF[4], float out[4]) const { - if (!computeBayesianPid.value || !hasTpc) { + if (!computeBayesianPid.value || !hasTpcIn) { out[0] = out[1] = out[2] = out[3] = kNaN; return; } @@ -245,7 +301,7 @@ struct PidFeatureExtractor { float sum = 0.f; for (int i = 0; i < 4; i++) { float logL = -0.5f * nsTPC[i] * nsTPC[i]; - if (hasTof) { + if (hasTofIn) { logL += -0.5f * nsTOF[i] * nsTOF[i]; } out[i] = std::exp(logL) * priors[i]; @@ -268,96 +324,119 @@ struct PidFeatureExtractor { 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 - struct HmpidRow { - bool has = false; - float signal = kNaN, qMip = kNaN, mom = kNaN; - int nPhotons = 0, clusSize = 0; - - static HmpidRow lookup(TTrack const& track, std::unordered_map const& map) - { - HmpidRow row; - if (auto it = map.find(track.globalIndex()); it != map.end()) { - row.has = true; - row.signal = it->second.hmpidSignal(); - row.qMip = it->second.hmpidQMip(); - row.mom = it->second.hmpidMom(); - row.nPhotons = it->second.hmpidNPhotons(); - row.clusSize = it->second.hmpidClusSize(); - } - return row; + 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() > -900.f; + 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(); } - }; - /// Builds the shared reconstructed-feature row for one track. Identical - /// for MC and data - the only thing that differs between the two modes is - /// whether MC-truth columns get appended afterwards. - template - FeatureRow buildFeatureRow(TTrack const& track, float vz, float centFT0C, - std::unordered_map const& hmpidMap) const - { - FeatureRow r{}; - r.p = track.p(); - r.pt = track.pt(); - r.px = track.px(); - r.py = track.py(); - r.pz = track.pz(); - r.eta = track.eta(); - r.phi = track.phi(); - r.sign = static_cast(track.sign()); - r.trackType = track.trackType(); - r.vz = vz; - r.centFT0C = centFT0C; - r.dcaXY = track.dcaXY(); - r.dcaZ = track.dcaZ(); - - r.hasTpc = track.hasTPC(); - r.tpcSignal = track.tpcSignal(); - r.tpcNSigmaPi = track.tpcNSigmaPi(); - r.tpcNSigmaKa = track.tpcNSigmaKa(); - r.tpcNSigmaPr = track.tpcNSigmaPr(); - r.tpcNSigmaEl = track.tpcNSigmaEl(); - r.tpcNClsFound = track.tpcNClsFound(); - r.tpcChi2NCl = track.tpcChi2NCl(); - - r.hasTof = !tofMissing(track); - r.tofMass = getTofMass(track); - r.beta = track.beta(); - r.tofNSigmaPi = track.tofNSigmaPi(); - r.tofNSigmaKa = track.tofNSigmaKa(); - r.tofNSigmaPr = track.tofNSigmaPr(); - r.tofNSigmaEl = track.tofNSigmaEl(); - - r.hasTrd = !trdMissing(track); - r.trdSignal = track.trdSignal(); - r.trdChi2 = track.trdChi2(); - r.trdPattern = track.trdPattern(); - - r.itsClusterSizes = track.itsClusterSizes(); - r.itsChi2NCl = track.itsChi2NCl(); - - r.hasEmcal = track.trackEtaEmcal() > -900.f; - r.trackEtaEmcal = track.trackEtaEmcal(); - r.trackPhiEmcal = track.trackPhiEmcal(); - - auto hm = HmpidRow::lookup(track, hmpidMap); - r.hasHmpid = hm.has; - r.hmpidSignal = hm.signal; - r.hmpidQMip = hm.qMip; - r.hmpidNPhotons = hm.nPhotons; - r.hmpidClusSize = hm.clusSize; - r.hmpidMom = hm.mom; - - float nsTPC[4] = {r.tpcNSigmaPi, r.tpcNSigmaKa, r.tpcNSigmaPr, r.tpcNSigmaEl}; - float nsTOF[4] = {r.tofNSigmaPi, r.tofNSigmaKa, r.tofNSigmaPr, r.tofNSigmaEl}; + float nsTPC[4] = {tpcNSigmaPi, tpcNSigmaKa, tpcNSigmaPr, tpcNSigmaEl}; + float nsTOF[4] = {tofNSigmaPi, tofNSigmaKa, tofNSigmaPr, tofNSigmaEl}; float bayes[4]; - computeBayesianProbs(r.hasTpc, nsTPC, r.hasTof, nsTOF, bayes); - r.bayesProbPi = bayes[0]; - r.bayesProbKa = bayes[1]; - r.bayesProbPr = bayes[2]; - r.bayesProbEl = bayes[3]; + computeBayesianProbs(hasTpc, nsTPC, hasTof, nsTOF, bayes); + bayesProbPi = bayes[0]; + bayesProbKa = bayes[1]; + bayesProbPr = bayes[2]; + bayesProbEl = bayes[3]; + } - return r; + 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) @@ -367,23 +446,9 @@ struct PidFeatureExtractor { if (!passesDpgCuts(track)) { continue; } - auto r = buildFeatureRow(track, collision.posZ(), collision.centFT0C(), hmpidMap); - - pidFeaturesData( - r.p, r.pt, r.px, r.py, r.pz, r.eta, r.phi, r.sign, r.trackType, - r.vz, r.centFT0C, r.dcaXY, r.dcaZ, - r.hasTpc, r.tpcSignal, r.tpcNSigmaPi, r.tpcNSigmaKa, r.tpcNSigmaPr, r.tpcNSigmaEl, - r.tpcNClsFound, r.tpcChi2NCl, - r.hasTof, r.tofMass, r.beta, r.tofNSigmaPi, r.tofNSigmaKa, r.tofNSigmaPr, r.tofNSigmaEl, - r.hasTrd, r.trdSignal, r.trdChi2, r.trdPattern, - r.itsClusterSizes, r.itsChi2NCl, - r.hasEmcal, r.trackEtaEmcal, r.trackPhiEmcal, - r.hasHmpid, r.hmpidSignal, r.hmpidQMip, r.hmpidNPhotons, r.hmpidClusSize, r.hmpidMom, - r.bayesProbPi, r.bayesProbKa, r.bayesProbPr, r.bayesProbEl); - - if (exportCsv.value) { - writeCsvRow(*csvFile, r); - } + fillRow(track, collision.posZ(), collision.centFT0C(), hmpidMap); + fillOutputs(); + fillQcHistos(); } } PROCESS_SWITCH(PidFeatureExtractor, processData, "Produce PID features for real/raw data (no MC truth)", true); @@ -399,27 +464,26 @@ struct PidFeatureExtractor { continue; } auto mcParticle = track.mcParticle(); - auto r = buildFeatureRow(track, collision.posZ(), collision.centFT0C(), hmpidMap); - - pidFeaturesMc( - r.p, r.pt, r.px, r.py, r.pz, r.eta, r.phi, r.sign, r.trackType, - r.vz, r.centFT0C, r.dcaXY, r.dcaZ, - r.hasTpc, r.tpcSignal, r.tpcNSigmaPi, r.tpcNSigmaKa, r.tpcNSigmaPr, r.tpcNSigmaEl, - r.tpcNClsFound, r.tpcChi2NCl, - r.hasTof, r.tofMass, r.beta, r.tofNSigmaPi, r.tofNSigmaKa, r.tofNSigmaPr, r.tofNSigmaEl, - r.hasTrd, r.trdSignal, r.trdChi2, r.trdPattern, - r.itsClusterSizes, r.itsChi2NCl, - r.hasEmcal, r.trackEtaEmcal, r.trackPhiEmcal, - r.hasHmpid, r.hmpidSignal, r.hmpidQMip, r.hmpidNPhotons, r.hmpidClusSize, r.hmpidMom, - r.bayesProbPi, r.bayesProbKa, r.bayesProbPr, r.bayesProbEl, - mcParticle.pdgCode(), static_cast(mcParticle.isPhysicalPrimary())); - - if (exportCsv.value) { - writeCsvRow(*csvFile, r); - } + 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) From 3439a14c9b728feebffb6081f541c416834c47af Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 17:20:38 +0100 Subject: [PATCH 08/16] Temporarily disable pid-onnx-inference pending rework for TTree-based extractor output --- Tools/PIDFeatureExtractor/CMakeLists.txt | 23 ++++++++--------------- 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/Tools/PIDFeatureExtractor/CMakeLists.txt b/Tools/PIDFeatureExtractor/CMakeLists.txt index 00ec05462cc..f4f544a5ad4 100644 --- a/Tools/PIDFeatureExtractor/CMakeLists.txt +++ b/Tools/PIDFeatureExtractor/CMakeLists.txt @@ -1,20 +1,13 @@ -# 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) \ No newline at end of file +# pid-onnx-inference temporarily disabled - it consumed aod::PidFeaturesData/ +# aod::PidFeaturesMc, which no longer exist now that pidFeatureExtractor.cxx +# writes TFile/TTree instead of O2 AOD tables. Needs rework before +# re-enabling - see project notes. +# o2physics_add_dpl_workflow(pid-onnx-inference +# SOURCES pidOnnxInference.cxx +# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore +# COMPONENT_NAME Analysis) From f7e15d14b7b6f9c9c73e326ae8d32a5498d84441 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 17:30:54 +0100 Subject: [PATCH 09/16] Rewrite pidOnnxInference as a one-shot batch task reading the extractor's TTree directly --- Tools/PIDFeatureExtractor/CMakeLists.txt | 12 +- .../PIDFeatureExtractor/pidOnnxInference.cxx | 393 ++++++++++-------- 2 files changed, 233 insertions(+), 172 deletions(-) diff --git a/Tools/PIDFeatureExtractor/CMakeLists.txt b/Tools/PIDFeatureExtractor/CMakeLists.txt index f4f544a5ad4..f5520cea85d 100644 --- a/Tools/PIDFeatureExtractor/CMakeLists.txt +++ b/Tools/PIDFeatureExtractor/CMakeLists.txt @@ -3,11 +3,7 @@ o2physics_add_dpl_workflow(pid-feature-extractor PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore COMPONENT_NAME Analysis) -# pid-onnx-inference temporarily disabled - it consumed aod::PidFeaturesData/ -# aod::PidFeaturesMc, which no longer exist now that pidFeatureExtractor.cxx -# writes TFile/TTree instead of O2 AOD tables. Needs rework before -# re-enabling - see project notes. -# o2physics_add_dpl_workflow(pid-onnx-inference -# SOURCES pidOnnxInference.cxx -# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore -# 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/pidOnnxInference.cxx b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx index fa64d223cf2..0448d1a1130 100644 --- a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx +++ b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx @@ -11,86 +11,49 @@ /// \file pidOnnxInference.cxx /// \brief Run the FSE PID ONNX model (loaded from CCDB, or a local file for -/// testing) over the tables produced by pidFeatureExtractor.cxx, and -/// write out per-track class probabilities. +/// testing) over the ROOT TTree produced by pidFeatureExtractor.cxx, +/// and write per-track class probabilities to a new ROOT file +/// (and/or CSV). /// -/// Uses Tools/ML/MlResponse - O2Physics's generic ONNX/CCDB inference -/// wrapper (not PID-specific, used across several PWG groups) - -/// rather than hand-rolled ONNX/CCDB loading, so this task stays -/// small. It depends only on that shared ML infrastructure and on -/// pidFeatureExtractor.h in this same folder - no other analysis -/// task. +/// Deliberately NOT an AOD-table-subscribing DPL task: it reads the +/// input file directly via plain TFile/TTree in init(), same as +/// pidFeatureExtractor.cxx now writes its output - no +/// DECLARE_SOA_TABLE, no Produces<>, avoiding the framework issue +/// that broke the earlier table-based version of this pair of +/// tasks tonight. Uses o2::analysis::MlResponse +/// (Tools/ML/MlResponse.h) - O2Physics's generic ONNX/CCDB +/// inference wrapper - for model loading and execution. /// /// \author Robert Forynski -#include "pidFeatureExtractor.h" -// #include "Tools/ML/MlResponse.h" -#include #include #include +#include #include +#include +#include + #include +#include +#include #include #include using namespace o2; using namespace o2::analysis; using namespace o2::framework; -using namespace o2::framework::expressions; - -// ============================================================================ -// OUTPUT TABLE -// ---------------------------------------------------------------------------- -// Declared directly here since nothing else needs to depend on it. -// ============================================================================ -namespace o2::aod -{ -namespace pidpred -{ -DECLARE_SOA_COLUMN(MlProbPi, mlProbPi, float); //! -DECLARE_SOA_COLUMN(MlProbKa, mlProbKa, float); //! -DECLARE_SOA_COLUMN(MlProbPr, mlProbPr, float); //! -DECLARE_SOA_COLUMN(MlProbEl, mlProbEl, float); //! -DECLARE_SOA_COLUMN(MlPredictedClass, mlPredictedClass, int); //! argmax of the 4 probs above: 0=pi,1=ka,2=pr,3=el -} // namespace pidpred - -DECLARE_SOA_TABLE(PidMlPredictions, "AOD", "PIDMLPRED", //! - o2::soa::Index<>, - pidpred::MlProbPi, pidpred::MlProbKa, pidpred::MlProbPr, pidpred::MlProbEl, - pidpred::MlPredictedClass); -} // namespace o2::aod namespace { constexpr int kNumClasses = 4; // pi, ka, pr, el - fixed order throughout, matches the paper's model -// ---------------------------------------------------------------------------- -// Feature order fed to the model. -// ---------------------------------------------------------------------------- -// THIS MUST MATCH YOUR TRAINING SCRIPT'S COLUMN ORDER EXACTLY - a silent -// mismatch here is the single most likely way this task produces wrong -// predictions without any error. The list below is a reasonable default -// (every reconstructed feature in PidFeaturesData/PidFeaturesMc except -// vz/centFT0C/sign/trackType and the Bayesian columns, which are a -// comparison baseline, not a model input) - it is NOT verified against your -// actual training code. Reorder, add, or drop entries to match exactly -// before trusting the output. -// -// If your ONNX export takes features and mask as two separate input -// tensors rather than one concatenated vector, split this function -// accordingly. - /// 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. -/// Small and duplicated here rather than shared with pidFeatureExtractor.cxx -/// so this task stays self-contained. -template -int getItsNClusters(TRow const& row) +int getItsNClusters(uint32_t v) { - auto v = static_cast(row.itsClusterSizes()); int n = 0; for (int layer = 0; layer < 7; layer++) { if ((v >> (layer * 4)) & 0xF) { @@ -100,75 +63,6 @@ int getItsNClusters(TRow const& row) return n; } -template -std::vector buildModelInput(TRow const& row) -{ - std::vector x; - x.reserve(38 + 7); - - // Kinematics - x.push_back(row.p()); - x.push_back(row.pt()); - x.push_back(row.px()); - x.push_back(row.py()); - x.push_back(row.pz()); - x.push_back(row.eta()); - x.push_back(row.phi()); - // Impact parameters - x.push_back(row.dcaXY()); - x.push_back(row.dcaZ()); - // TPC - x.push_back(static_cast(row.hasTPC())); - x.push_back(row.tpcSignal()); - x.push_back(row.tpcNSigmaPi()); - x.push_back(row.tpcNSigmaKa()); - x.push_back(row.tpcNSigmaPr()); - x.push_back(row.tpcNSigmaEl()); - x.push_back(static_cast(row.tpcNClsFound())); - x.push_back(row.tpcChi2NCl()); - // TOF - x.push_back(static_cast(row.hasTOF())); - x.push_back(row.tofMass()); - x.push_back(row.beta()); - x.push_back(row.tofNSigmaPi()); - x.push_back(row.tofNSigmaKa()); - x.push_back(row.tofNSigmaPr()); - x.push_back(row.tofNSigmaEl()); - // TRD - x.push_back(static_cast(row.hasTRD())); - x.push_back(row.trdSignal()); - x.push_back(row.trdChi2()); - x.push_back(static_cast(row.trdPattern())); - // ITS - x.push_back(static_cast(getItsNClusters(row))); - x.push_back(row.itsChi2NCl()); - // EMCal - x.push_back(static_cast(row.hasEMCal())); - x.push_back(row.trackEtaEmcal()); - x.push_back(row.trackPhiEmcal()); - // HMPID - x.push_back(static_cast(row.hasHMPID())); - x.push_back(row.hmpidSignal()); - x.push_back(row.hmpidQMip()); - x.push_back(static_cast(row.hmpidNPhotons())); - x.push_back(static_cast(row.hmpidClusSize())); - x.push_back(row.hmpidMom()); - - // 7-length group mask: TPC, TOF, TRD, ITS, EMCal, HMPID, centrality. - // ITS is assumed always present (global track requires it); centrality is - // assumed always present. Both are real assumptions, not derived facts - - // adjust if your training data ever has either group absent. - x.push_back(static_cast(row.hasTPC())); - x.push_back(static_cast(row.hasTOF())); - x.push_back(static_cast(row.hasTRD())); - x.push_back(1.f); // ITS - x.push_back(static_cast(row.hasEMCal())); - x.push_back(static_cast(row.hasHMPID())); - x.push_back(1.f); // centrality - - return x; -} - int argmax4(std::vector const& v) { int best = 0; @@ -181,24 +75,35 @@ int argmax4(std::vector const& v) } } // namespace -/// PidOnnxInference: applies the FSE ONNX model to features produced by +/// PidOnnxInference: applies the FSE ONNX model to the features written by /// pidFeatureExtractor.cxx and writes out per-track class probabilities. /// -/// Model loading (local file or CCDB) and ONNX execution are entirely -/// handled by o2::analysis::MlResponse - this task only builds the input -/// feature vector and reads back the output. A single global pT bin is used -/// by default since the paper's model isn't pT-binned; MlResponse's usual -/// per-class cut mechanism is disabled (cutDirMl = CutNot for every class, -/// confirmed value 2 in the real o2::cuts_ml enum) - this task always -/// reports all four probabilities rather than applying a pass/fail cut. +/// This is a one-shot batch task, not a per-timeframe DPL processor: all +/// real work happens in init() (load model, read the whole input tree, +/// write predictions), since the input is a finished file from a previous +/// run rather than live AO2D data. A trivial empty process() is kept so +/// the task still registers normally with adaptAnalysisTask, matching +/// every other task in this project - NOT VERIFIED that +/// ControlService::endOfStream()/readyToQuit() is the correct way to make +/// the workflow terminate cleanly after init() finishes; if the workflow +/// hangs instead of exiting, this is the first thing to revisit. /// -/// Two things MlResponse enforces that are worth knowing before debugging a -/// failure here: getModelOutput() calls LOG(fatal) if the input vector's -/// length doesn't match the ONNX model's declared input node count (unless -/// that node is a dynamic axis), and separately if pt lands outside -/// binsPtMl's range entirely (see the comment on binsPtMl below). +/// - cutDirMl defaults to cuts_ml::CutNot (value 2) for every class, so +/// MlResponse never rejects a track - this task always reports all four +/// probabilities rather than making a pass/fail decision. +/// - A single pT bin is used by default (model isn't pT-binned); lower +/// edge is -1, not 0, so no track's pt() can land exactly on the +/// boundary (MlResponse::findBin() treats that as out-of-range). +/// - buildModelInput()'s feature order is a REASONABLE DEFAULT, not +/// verified against the actual training code - see the comment there. struct PidOnnxInference { - Produces pidMlPredictions; + // --- input: the ROOT file/tree pidFeatureExtractor.cxx wrote ------------- + Configurable inputRootFile{"inputRootFile", "pid_features_data.root", "ROOT file produced by pidFeatureExtractor.cxx"}; + Configurable inputTreeName{"inputTreeName", "pid_features", "Name of the TTree inside inputRootFile"}; + + // --- output ------------------------------------------------------------------ + Configurable outputPath{"outputPath", "pid_predictions", "Output file base name (no extension)"}; + Configurable exportCsv{"exportCsv", false, "Also write predictions to CSV alongside the ROOT output"}; // --- model location ---------------------------------------------------------- Configurable loadModelFromCcdb{"loadModelFromCcdb", true, "Load the ONNX model from CCDB (else from onnxFileNames as a local path)"}; @@ -208,12 +113,7 @@ struct PidOnnxInference { Configurable> onnxFileNames{"onnxFileNames", std::vector{"pid_feature_model.onnx"}, "Local ONNX file path(s), used when loadModelFromCcdb is false"}; // --- MlResponse plumbing: a single pT bin, no selection cut applied -------- - // Lower edge is -1 (not 0): MlResponse::findBin() rejects value < front() - // as out-of-range (fatal in getModelOutput), and track.pt() could in - // principle be exactly 0 - keeping the edge below any physical pT avoids - // that boundary case entirely. Configurable> binsPtMl{"binsPtMl", std::vector{-1., 9999.}, "pT bin edges for MlResponse (single bin = model isn't pT-binned)"}; - Configurable> cutDirMl{"cutDirMl", std::vector{cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot}, "Per-class cut direction; CutNot = always accept, this task doesn't select"}; static constexpr double kDefaultCutsMl[1][kNumClasses] = {{0., 0., 0., 0.}}; Configurable> cutsMl{"cutsMl", {kDefaultCutsMl[0], 1, kNumClasses, {"pT bin 0"}, {"prob pi", "prob ka", "prob pr", "prob el"}}, "Unused thresholds (CutNot everywhere) - required by MlResponse's interface"}; @@ -221,42 +121,207 @@ struct PidOnnxInference { o2::ccdb::CcdbApi ccdbApi; o2::analysis::MlResponse mlResponse; - std::vector mlOutput; - void init(InitContext const&) + void init(InitContext& ic) { - mlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); - if (loadModelFromCcdb) { + mlResponse.configure(binsPtMl, cutsMl, std::vector{cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot}, nClassesMl); + if (loadModelFromCcdb.value) { ccdbApi.init(ccdbUrl.value); mlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCcdb.value, timestampCcdb.value); } else { mlResponse.setModelPathsLocal(onnxFileNames); } mlResponse.init(); + + runInferenceOverFile(); + + // One-shot batch job: nothing to subscribe to, so tell DPL we're done + // rather than waiting for input that will never arrive. + ic.services().get().endOfStream(); + ic.services().get().readyToQuit(QuitRequest::Me); } - template - void runInference(TTable const& rows) + /// Trivial, intentionally empty - exists only so this task registers + /// like every other task in the project. All real work is in init(). + void process(ProcessingContext&) {} + + void runInferenceOverFile() { - for (auto const& row : rows) { - auto x = buildModelInput(row); - mlResponse.isSelectedMl(x, row.pt(), mlOutput); // return value (selection) unused; mlOutput carries the 4 raw scores - pidMlPredictions(mlOutput[0], mlOutput[1], mlOutput[2], mlOutput[3], argmax4(mlOutput)); - mlOutput.clear(); + std::unique_ptr inFile(TFile::Open(inputRootFile.value.c_str(), "READ")); + if (!inFile || inFile->IsZombie()) { + LOG(fatal) << "Could not open input file " << inputRootFile.value; + return; + } + auto* tree = dynamic_cast(inFile->Get(inputTreeName.value.c_str())); + if (!tree) { + LOG(fatal) << "Tree " << inputTreeName.value << " not found in " << inputRootFile.value; + return; } - } - void processData(aod::PidFeaturesData const& rows) - { - runInference(rows); - } - PROCESS_SWITCH(PidOnnxInference, processData, "Run inference on PidFeaturesData", true); + // 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; - void processMc(aod::PidFeaturesMc const& rows) - { - runInference(rows); + 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.value + ".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.value) { + csv.open(outputPath.value + ".csv"); + csv << "mlProbPi,mlProbKa,mlProbPr,mlProbEl,mlPredictedClass\n"; + } + + // -------------------------------------------------------------------------- + // 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; ITS and centrality are + // assumed always-present). NOT verified against your actual training code. + // -------------------------------------------------------------------------- + std::vector x; + std::vector mlOutput; + Long64_t nEntries = tree->GetEntries(); + for (Long64_t i = 0; i < nEntries; i++) { + tree->GetEntry(i); + + x.clear(); + x.reserve(38 + 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(hasTPC)); + x.push_back(tpcSignal); + x.push_back(tpcNSigmaPi); + x.push_back(tpcNSigmaKa); + x.push_back(tpcNSigmaPr); + x.push_back(tpcNSigmaEl); + x.push_back(static_cast(tpcNClsFound)); + x.push_back(tpcChi2NCl); + x.push_back(static_cast(hasTOF)); + x.push_back(tofMass); + x.push_back(beta); + x.push_back(tofNSigmaPi); + x.push_back(tofNSigmaKa); + x.push_back(tofNSigmaPr); + x.push_back(tofNSigmaEl); + x.push_back(static_cast(hasTRD)); + x.push_back(trdSignal); + x.push_back(trdChi2); + x.push_back(static_cast(trdPattern)); + x.push_back(static_cast(getItsNClusters(static_cast(itsClusterSizes)))); + x.push_back(itsChi2NCl); + x.push_back(static_cast(hasEMCal)); + x.push_back(trackEtaEmcal); + x.push_back(trackPhiEmcal); + x.push_back(static_cast(hasHMPID)); + x.push_back(hmpidSignal); + x.push_back(hmpidQMip); + x.push_back(static_cast(hmpidNPhotons)); + x.push_back(static_cast(hmpidClusSize)); + x.push_back(hmpidMom); + // 7-length group mask + x.push_back(static_cast(hasTPC)); + x.push_back(static_cast(hasTOF)); + x.push_back(static_cast(hasTRD)); + x.push_back(1.f); // ITS + x.push_back(static_cast(hasEMCal)); + x.push_back(static_cast(hasHMPID)); + x.push_back(1.f); // centrality + + 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.value) { + csv << mlProbPi << ',' << mlProbKa << ',' << mlProbPr << ',' << mlProbEl << ',' << mlPredictedClass << '\n'; + } + } + + outFile->cd(); + outTree.Write(); + outFile->Close(); + if (exportCsv.value) { + csv.close(); + } + + LOG(info) << "PidOnnxInference: wrote " << nEntries << " predictions to " << outputPath.value << ".root"; } - PROCESS_SWITCH(PidOnnxInference, processMc, "Run inference on PidFeaturesMc", false); }; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) From f21f33e5f6c08b6857b4c560f8f967de676e2360 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 17:35:38 +0100 Subject: [PATCH 10/16] Rewrite pidOnnxInference as a plain DataProcessorSpec, avoiding adaptAnalysisTask's AOD-table-only process() reflection --- .../PIDFeatureExtractor/pidOnnxInference.cxx | 487 +++++++++--------- 1 file changed, 243 insertions(+), 244 deletions(-) diff --git a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx index 0448d1a1130..c37657000c9 100644 --- a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx +++ b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx @@ -15,22 +15,29 @@ /// and write per-track class probabilities to a new ROOT file /// (and/or CSV). /// -/// Deliberately NOT an AOD-table-subscribing DPL task: it reads the -/// input file directly via plain TFile/TTree in init(), same as -/// pidFeatureExtractor.cxx now writes its output - no -/// DECLARE_SOA_TABLE, no Produces<>, avoiding the framework issue -/// that broke the earlier table-based version of this pair of -/// tasks tonight. Uses o2::analysis::MlResponse -/// (Tools/ML/MlResponse.h) - O2Physics's generic ONNX/CCDB -/// inference wrapper - for model loading and execution. +/// Built as a plain DataProcessorSpec/AlgorithmSpec, NOT +/// adaptAnalysisTask<> - this is a standalone batch job with no AOD +/// table subscription, and adaptAnalysisTask's process() reflection +/// specifically requires an AOD table/iterator argument (confirmed +/// by a real compiler error: a raw ProcessingContext& process() +/// signature is rejected outright). Using the lower-level DPL +/// primitives directly avoids that machinery entirely, rather than +/// trying to satisfy a reflection mechanism built for a different +/// kind of task. Configurable<> member auto-binding is an +/// AnalysisTask-specific convenience, so options here are declared +/// explicitly and read via ic.options().get(...) instead. +/// +/// 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 @@ -73,258 +80,250 @@ int argmax4(std::vector const& v) } return best; } -} // namespace -/// PidOnnxInference: applies the FSE ONNX model to the features written by -/// pidFeatureExtractor.cxx and writes out per-track class probabilities. +/// All the real work: load the model, read the whole input tree, run +/// inference row by row, write predictions. Called once from init(). /// -/// This is a one-shot batch task, not a per-timeframe DPL processor: all -/// real work happens in init() (load model, read the whole input tree, -/// write predictions), since the input is a finished file from a previous -/// run rather than live AO2D data. A trivial empty process() is kept so -/// the task still registers normally with adaptAnalysisTask, matching -/// every other task in this project - NOT VERIFIED that -/// ControlService::endOfStream()/readyToQuit() is the correct way to make -/// the workflow terminate cleanly after init() finishes; if the workflow -/// hangs instead of exiting, this is the first thing to revisit. -/// -/// - cutDirMl defaults to cuts_ml::CutNot (value 2) for every class, so -/// MlResponse never rejects a track - this task always reports all four -/// probabilities rather than making a pass/fail decision. -/// - A single pT bin is used by default (model isn't pT-binned); lower -/// edge is -1, not 0, so no track's pt() can land exactly on the -/// boundary (MlResponse::findBin() treats that as out-of-range). -/// - buildModelInput()'s feature order is a REASONABLE DEFAULT, not -/// verified against the actual training code - see the comment there. -struct PidOnnxInference { - // --- input: the ROOT file/tree pidFeatureExtractor.cxx wrote ------------- - Configurable inputRootFile{"inputRootFile", "pid_features_data.root", "ROOT file produced by pidFeatureExtractor.cxx"}; - Configurable inputTreeName{"inputTreeName", "pid_features", "Name of the TTree inside inputRootFile"}; - - // --- output ------------------------------------------------------------------ - Configurable outputPath{"outputPath", "pid_predictions", "Output file base name (no extension)"}; - Configurable exportCsv{"exportCsv", false, "Also write predictions to CSV alongside the ROOT output"}; - - // --- model location ---------------------------------------------------------- - 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"}; - - // --- MlResponse plumbing: a single pT bin, no selection cut applied -------- - Configurable> binsPtMl{"binsPtMl", std::vector{-1., 9999.}, "pT bin edges for MlResponse (single bin = model isn't pT-binned)"}; - - static constexpr double kDefaultCutsMl[1][kNumClasses] = {{0., 0., 0., 0.}}; - Configurable> cutsMl{"cutsMl", {kDefaultCutsMl[0], 1, kNumClasses, {"pT bin 0"}, {"prob pi", "prob ka", "prob pr", "prob el"}}, "Unused thresholds (CutNot everywhere) - required by MlResponse's interface"}; - Configurable nClassesMl{"nClassesMl", static_cast(kNumClasses), "Number of model output classes"}; +/// 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; ITS and centrality are +/// assumed always-present). NOT verified against your actual training code. +void runInference(std::string const& inputRootFile, std::string const& inputTreeName, + std::string const& outputPath, bool exportCsv, + o2::analysis::MlResponse& mlResponse) +{ + 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; + } - o2::ccdb::CcdbApi ccdbApi; - o2::analysis::MlResponse mlResponse; + // 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; - void init(InitContext& ic) - { - mlResponse.configure(binsPtMl, cutsMl, std::vector{cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot, cuts_ml::CutNot}, nClassesMl); - if (loadModelFromCcdb.value) { - ccdbApi.init(ccdbUrl.value); - mlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCcdb.value, timestampCcdb.value); - } else { - mlResponse.setModelPathsLocal(onnxFileNames); - } - mlResponse.init(); + 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); - runInferenceOverFile(); + 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); - // One-shot batch job: nothing to subscribe to, so tell DPL we're done - // rather than waiting for input that will never arrive. - ic.services().get().endOfStream(); - ic.services().get().readyToQuit(QuitRequest::Me); + std::ofstream csv; + if (exportCsv) { + csv.open(outputPath + ".csv"); + csv << "mlProbPi,mlProbKa,mlProbPr,mlProbEl,mlPredictedClass\n"; } - /// Trivial, intentionally empty - exists only so this task registers - /// like every other task in the project. All real work is in init(). - void process(ProcessingContext&) {} + std::vector x; + std::vector mlOutput; + Long64_t nEntries = tree->GetEntries(); + for (Long64_t i = 0; i < nEntries; i++) { + tree->GetEntry(i); - void runInferenceOverFile() - { - std::unique_ptr inFile(TFile::Open(inputRootFile.value.c_str(), "READ")); - if (!inFile || inFile->IsZombie()) { - LOG(fatal) << "Could not open input file " << inputRootFile.value; - return; - } - auto* tree = dynamic_cast(inFile->Get(inputTreeName.value.c_str())); - if (!tree) { - LOG(fatal) << "Tree " << inputTreeName.value << " not found in " << inputRootFile.value; - return; - } + x.clear(); + x.reserve(38 + 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(hasTPC)); + x.push_back(tpcSignal); + x.push_back(tpcNSigmaPi); + x.push_back(tpcNSigmaKa); + x.push_back(tpcNSigmaPr); + x.push_back(tpcNSigmaEl); + x.push_back(static_cast(tpcNClsFound)); + x.push_back(tpcChi2NCl); + x.push_back(static_cast(hasTOF)); + x.push_back(tofMass); + x.push_back(beta); + x.push_back(tofNSigmaPi); + x.push_back(tofNSigmaKa); + x.push_back(tofNSigmaPr); + x.push_back(tofNSigmaEl); + x.push_back(static_cast(hasTRD)); + x.push_back(trdSignal); + x.push_back(trdChi2); + x.push_back(static_cast(trdPattern)); + x.push_back(static_cast(getItsNClusters(static_cast(itsClusterSizes)))); + x.push_back(itsChi2NCl); + x.push_back(static_cast(hasEMCal)); + x.push_back(trackEtaEmcal); + x.push_back(trackPhiEmcal); + x.push_back(static_cast(hasHMPID)); + x.push_back(hmpidSignal); + x.push_back(hmpidQMip); + x.push_back(static_cast(hmpidNPhotons)); + x.push_back(static_cast(hmpidClusSize)); + x.push_back(hmpidMom); + // 7-length group mask + x.push_back(static_cast(hasTPC)); + x.push_back(static_cast(hasTOF)); + x.push_back(static_cast(hasTRD)); + x.push_back(1.f); // ITS + x.push_back(static_cast(hasEMCal)); + x.push_back(static_cast(hasHMPID)); + x.push_back(1.f); // centrality - // 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; + 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(); - 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.value + ".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.value) { - csv.open(outputPath.value + ".csv"); - csv << "mlProbPi,mlProbKa,mlProbPr,mlProbEl,mlPredictedClass\n"; + if (exportCsv) { + csv << mlProbPi << ',' << mlProbKa << ',' << mlProbPr << ',' << mlProbEl << ',' << mlPredictedClass << '\n'; } + } - // -------------------------------------------------------------------------- - // 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; ITS and centrality are - // assumed always-present). NOT verified against your actual training code. - // -------------------------------------------------------------------------- - std::vector x; - std::vector mlOutput; - Long64_t nEntries = tree->GetEntries(); - for (Long64_t i = 0; i < nEntries; i++) { - tree->GetEntry(i); + outFile->cd(); + outTree.Write(); + outFile->Close(); + if (exportCsv) { + csv.close(); + } + + LOG(info) << "PidOnnxInference: wrote " << nEntries << " predictions to " << outputPath << ".root"; +} +} // namespace - x.clear(); - x.reserve(38 + 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(hasTPC)); - x.push_back(tpcSignal); - x.push_back(tpcNSigmaPi); - x.push_back(tpcNSigmaKa); - x.push_back(tpcNSigmaPr); - x.push_back(tpcNSigmaEl); - x.push_back(static_cast(tpcNClsFound)); - x.push_back(tpcChi2NCl); - x.push_back(static_cast(hasTOF)); - x.push_back(tofMass); - x.push_back(beta); - x.push_back(tofNSigmaPi); - x.push_back(tofNSigmaKa); - x.push_back(tofNSigmaPr); - x.push_back(tofNSigmaEl); - x.push_back(static_cast(hasTRD)); - x.push_back(trdSignal); - x.push_back(trdChi2); - x.push_back(static_cast(trdPattern)); - x.push_back(static_cast(getItsNClusters(static_cast(itsClusterSizes)))); - x.push_back(itsChi2NCl); - x.push_back(static_cast(hasEMCal)); - x.push_back(trackEtaEmcal); - x.push_back(trackPhiEmcal); - x.push_back(static_cast(hasHMPID)); - x.push_back(hmpidSignal); - x.push_back(hmpidQMip); - x.push_back(static_cast(hmpidNPhotons)); - x.push_back(static_cast(hmpidClusSize)); - x.push_back(hmpidMom); - // 7-length group mask - x.push_back(static_cast(hasTPC)); - x.push_back(static_cast(hasTOF)); - x.push_back(static_cast(hasTRD)); - x.push_back(1.f); // ITS - x.push_back(static_cast(hasEMCal)); - x.push_back(static_cast(hasHMPID)); - x.push_back(1.f); // centrality +WorkflowSpec defineDataProcessing(ConfigContext const&) +{ + DataProcessorSpec spec{ + "pid-onnx-inference", + Inputs{}, + Outputs{}, + AlgorithmSpec{[](InitContext& ic) { + auto inputRootFile = ic.options().get("inputRootFile"); + auto inputTreeName = ic.options().get("inputTreeName"); + auto outputPath = ic.options().get("outputPath"); + auto exportCsv = ic.options().get("exportCsv"); + auto loadModelFromCcdb = ic.options().get("loadModelFromCcdb"); + auto ccdbUrl = ic.options().get("ccdbUrl"); + auto modelPathsCcdb = ic.options().get>("modelPathsCcdb"); + auto timestampCcdb = ic.options().get("timestampCcdb"); + auto onnxFileNames = ic.options().get>("onnxFileNames"); + auto binsPtMl = ic.options().get>("binsPtMl"); + auto nClassesMl = static_cast(ic.options().get("nClassesMl")); - 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(); + // 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; hardcoded here + // rather than exposed as options. + 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}; - if (exportCsv.value) { - csv << mlProbPi << ',' << mlProbKa << ',' << mlProbPr << ',' << mlProbEl << ',' << mlPredictedClass << '\n'; + auto mlResponse = std::make_shared>(); + mlResponse->configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); + if (loadModelFromCcdb) { + auto ccdbApi = std::make_shared(); + ccdbApi->init(ccdbUrl); + mlResponse->setModelPathsCCDB(onnxFileNames, *ccdbApi, modelPathsCcdb, timestampCcdb); + } else { + mlResponse->setModelPathsLocal(onnxFileNames); } - } + mlResponse->init(); - outFile->cd(); - outTree.Write(); - outFile->Close(); - if (exportCsv.value) { - csv.close(); - } + runInference(inputRootFile, inputTreeName, outputPath, exportCsv, *mlResponse); - LOG(info) << "PidOnnxInference: wrote " << nEntries << " predictions to " << outputPath.value << ".root"; - } -}; + return [](ProcessingContext& pc) { + // One-shot batch job: all work already happened in init(). Signal + // completion immediately rather than waiting for input that will + // never arrive. + pc.services().get().endOfStream(); + pc.services().get().readyToQuit(QuitRequest::Me); + }; + }}, + Options{ + {"inputRootFile", VariantType::String, "pid_features_data.root", {"ROOT file produced by pidFeatureExtractor.cxx"}}, + {"inputTreeName", VariantType::String, "pid_features", {"Name of the TTree inside inputRootFile"}}, + {"outputPath", VariantType::String, "pid_predictions", {"Output file base name (no extension)"}}, + {"exportCsv", VariantType::Bool, false, {"Also write predictions to CSV alongside the ROOT output"}}, + {"loadModelFromCcdb", VariantType::Bool, true, {"Load the ONNX model from CCDB (else from onnxFileNames as a local path)"}}, + {"ccdbUrl", VariantType::String, "http://alice-ccdb.cern.ch", {"CCDB URL"}}, + {"modelPathsCcdb", VariantType::ArrayString, std::vector{"Users/YOURNAME/PidFeatureExtractor/model"}, {"CCDB path to the model"}}, + {"timestampCcdb", VariantType::Int64, static_cast(-1), {"CCDB query timestamp for the model, -1 = latest"}}, + {"onnxFileNames", VariantType::ArrayString, std::vector{"pid_feature_model.onnx"}, {"Local ONNX file path(s), used when loadModelFromCcdb is false"}}, + {"binsPtMl", VariantType::ArrayDouble, std::vector{-1., 9999.}, {"pT bin edges for MlResponse (single bin = model isn't pT-binned)"}}, + {"nClassesMl", VariantType::Int, kNumClasses, {"Number of model output classes"}}, + }}; -WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) -{ - return WorkflowSpec{adaptAnalysisTask(cfgc)}; + return WorkflowSpec{spec}; } From 86aabca9a2ad583c10ecbf02503cb16696f06c64 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Fri, 7 Aug 2026 17:54:32 +0100 Subject: [PATCH 11/16] Add per-detector-group toggles to inference; rewrite README to match Tools/PIDML style --- Tools/PIDFeatureExtractor/README.md | 126 +++++++++++++++++ .../PIDFeatureExtractor/pidOnnxInference.cxx | 131 ++++++++++++------ 2 files changed, 215 insertions(+), 42 deletions(-) create mode 100644 Tools/PIDFeatureExtractor/README.md diff --git a/Tools/PIDFeatureExtractor/README.md b/Tools/PIDFeatureExtractor/README.md new file mode 100644 index 00000000000..a1f87849e76 --- /dev/null +++ b/Tools/PIDFeatureExtractor/README.md @@ -0,0 +1,126 @@ +# 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 - they're two +separate steps, not one pipeline (see "Running" below for why). + +## 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. + +### Configurable 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. The model can be loaded either from CCDB or from a +local file, which is handled by `o2::analysis::MlResponse` +(`Tools/ML/MlResponse.h`). + +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. + +### Configurable 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. + +`PidFeatureExtractor` needs to run as part of the normal AOD pipeline, +since it reads track and collision data directly: + +```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 +``` + +`PidOnnxInference` runs on its own, after that has finished - it just +opens the file the extractor wrote, so there's no AOD pipeline to build: + +```bash +#!/bin/bash + +config_file="my-config.json" + +o2-analysis-pid-onnx-inference --configuration json://$config_file -b +``` diff --git a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx index c37657000c9..5f3956dca39 100644 --- a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx +++ b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx @@ -45,6 +45,7 @@ #include #include +#include #include #include #include @@ -56,6 +57,7 @@ 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(); /// 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. @@ -81,6 +83,23 @@ int argmax4(std::vector const& v) 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(). /// @@ -88,11 +107,14 @@ int argmax4(std::vector const& v) /// 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; ITS and centrality are -/// assumed always-present). NOT verified against your actual training code. +/// 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) + o2::analysis::MlResponse& mlResponse, GroupToggles const& groups) { std::unique_ptr inFile(TFile::Open(inputRootFile.c_str(), "READ")); if (!inFile || inFile->IsZombie()) { @@ -189,8 +211,17 @@ void runInference(std::string const& inputRootFile, std::string const& inputTree for (Long64_t i = 0; i < nEntries; i++) { tree->GetEntry(i); + // Effective presence = what the data says AND the group is enabled. + // Disabling a group here forces the same "absent" state as a real + // detector miss, regardless of what hasXXX says in the input tree. + 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(38 + 7); + x.reserve(39 + 7); x.push_back(p); x.push_back(pt); x.push_back(px); @@ -200,44 +231,44 @@ void runInference(std::string const& inputRootFile, std::string const& inputTree x.push_back(phi); x.push_back(dcaXY); x.push_back(dcaZ); - x.push_back(static_cast(hasTPC)); - x.push_back(tpcSignal); - x.push_back(tpcNSigmaPi); - x.push_back(tpcNSigmaKa); - x.push_back(tpcNSigmaPr); - x.push_back(tpcNSigmaEl); - x.push_back(static_cast(tpcNClsFound)); - x.push_back(tpcChi2NCl); - x.push_back(static_cast(hasTOF)); - x.push_back(tofMass); - x.push_back(beta); - x.push_back(tofNSigmaPi); - x.push_back(tofNSigmaKa); - x.push_back(tofNSigmaPr); - x.push_back(tofNSigmaEl); - x.push_back(static_cast(hasTRD)); - x.push_back(trdSignal); - x.push_back(trdChi2); - x.push_back(static_cast(trdPattern)); - x.push_back(static_cast(getItsNClusters(static_cast(itsClusterSizes)))); - x.push_back(itsChi2NCl); - x.push_back(static_cast(hasEMCal)); - x.push_back(trackEtaEmcal); - x.push_back(trackPhiEmcal); - x.push_back(static_cast(hasHMPID)); - x.push_back(hmpidSignal); - x.push_back(hmpidQMip); - x.push_back(static_cast(hmpidNPhotons)); - x.push_back(static_cast(hmpidClusSize)); - x.push_back(hmpidMom); + 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(hasTPC)); - x.push_back(static_cast(hasTOF)); - x.push_back(static_cast(hasTRD)); - x.push_back(1.f); // ITS - x.push_back(static_cast(hasEMCal)); - x.push_back(static_cast(hasHMPID)); - x.push_back(1.f); // centrality + 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]; @@ -282,6 +313,15 @@ WorkflowSpec defineDataProcessing(ConfigContext const&) auto binsPtMl = ic.options().get>("binsPtMl"); auto nClassesMl = static_cast(ic.options().get("nClassesMl")); + GroupToggles groups; + groups.useTPC = ic.options().get("useTPC"); + groups.useTOF = ic.options().get("useTOF"); + groups.useTRD = ic.options().get("useTRD"); + groups.useITS = ic.options().get("useITS"); + groups.useEMCal = ic.options().get("useEMCal"); + groups.useHMPID = ic.options().get("useHMPID"); + groups.useCentrality = ic.options().get("useCentrality"); + // 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; hardcoded here @@ -301,7 +341,7 @@ WorkflowSpec defineDataProcessing(ConfigContext const&) } mlResponse->init(); - runInference(inputRootFile, inputTreeName, outputPath, exportCsv, *mlResponse); + runInference(inputRootFile, inputTreeName, outputPath, exportCsv, *mlResponse, groups); return [](ProcessingContext& pc) { // One-shot batch job: all work already happened in init(). Signal @@ -323,6 +363,13 @@ WorkflowSpec defineDataProcessing(ConfigContext const&) {"onnxFileNames", VariantType::ArrayString, std::vector{"pid_feature_model.onnx"}, {"Local ONNX file path(s), used when loadModelFromCcdb is false"}}, {"binsPtMl", VariantType::ArrayDouble, std::vector{-1., 9999.}, {"pT bin edges for MlResponse (single bin = model isn't pT-binned)"}}, {"nClassesMl", VariantType::Int, kNumClasses, {"Number of model output classes"}}, + {"useTPC", VariantType::Bool, true, {"Include TPC in inference. Default true (all detectors present); set false to force TPC excluded regardless of the data"}}, + {"useTOF", VariantType::Bool, true, {"Include TOF in inference"}}, + {"useTRD", VariantType::Bool, true, {"Include TRD in inference"}}, + {"useITS", VariantType::Bool, true, {"Include ITS in inference"}}, + {"useEMCal", VariantType::Bool, true, {"Include EMCal in inference"}}, + {"useHMPID", VariantType::Bool, true, {"Include HMPID in inference"}}, + {"useCentrality", VariantType::Bool, true, {"Include centrality in inference"}}, }}; return WorkflowSpec{spec}; From 5c8280b12066c6a252a460770f132f59c8e1ddcb Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Tue, 11 Aug 2026 21:15:11 +0100 Subject: [PATCH 12/16] Remove obsolete DataModel header, run.sh, and untrack local config JSON --- .../PIDFeatureExtractor/pidFeatureExtractor.h | 129 --------------- .../pidFeatureExtractorConfig.json | 150 ------------------ Tools/PIDFeatureExtractor/run.sh | 77 --------- 3 files changed, 356 deletions(-) delete mode 100644 Tools/PIDFeatureExtractor/pidFeatureExtractor.h delete mode 100644 Tools/PIDFeatureExtractor/pidFeatureExtractorConfig.json delete mode 100644 Tools/PIDFeatureExtractor/run.sh diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractor.h b/Tools/PIDFeatureExtractor/pidFeatureExtractor.h deleted file mode 100644 index 57ea3aae774..00000000000 --- a/Tools/PIDFeatureExtractor/pidFeatureExtractor.h +++ /dev/null @@ -1,129 +0,0 @@ -// 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.h -/// \brief Data model for the PID feature extractor: PidFeaturesData / -/// PidFeaturesMc table definitions. Included with a plain quoted -/// filename ("pidFeatureExtractor.h"), which the compiler resolves -/// relative to the including .cxx's own directory first - so this -/// header and every .cxx that uses it must stay in the same folder; -/// no repo-root-relative path to get wrong. -/// -/// \author Robert Forynski - -#ifndef PID_FEATURE_EXTRACTOR_H_ -#define PID_FEATURE_EXTRACTOR_H_ - -#include "Common/DataModel/Centrality.h" -#include "Common/DataModel/PIDResponseTOF.h" -#include "Common/DataModel/PIDResponseTPC.h" -#include "Common/DataModel/TrackSelectionTables.h" - -#include - -#include - -namespace o2::aod -{ -namespace pidfeat -{ -// Kinematics not already stored as static columns upstream (Pt, TrackType, -// DcaXY, DcaZ, TPCSignal, TRDSignal, TRDPattern, TrackEtaEMCAL, -// TrackPhiEMCAL, ITSClusterSizes, ITSChi2NCl, TPCNClsFound are reused -// directly from aod::track:: in the table definitions below). -DECLARE_SOA_COLUMN(P, p, float); //! Track momentum magnitude (GeV/c) -DECLARE_SOA_COLUMN(Px, px, float); //! Track x-momentum (GeV/c) -DECLARE_SOA_COLUMN(Py, py, float); //! Track y-momentum (GeV/c) -DECLARE_SOA_COLUMN(Pz, pz, float); //! Track z-momentum (GeV/c) -DECLARE_SOA_COLUMN(Eta, eta, float); //! Pseudorapidity -DECLARE_SOA_COLUMN(Phi, phi, float); //! Azimuthal angle -DECLARE_SOA_COLUMN(Sign, sign, float); //! Track charge sign - -// Event level, duplicated per track for a self-contained flat table. -DECLARE_SOA_COLUMN(Vz, vz, float); //! Collision vertex z (cm) - -// TOF mass: not exposed as a static column upstream. -DECLARE_SOA_COLUMN(TofMass, tofMass, float); //! TOF-reconstructed mass, NaN if !hasTOF() - -// Detector-presence flags. Kept explicit (rather than inferred solely from -// NaN sentinels) because the downstream FSE/attention model conditions on a -// per-track detector mask. -DECLARE_SOA_COLUMN(HasTPC, hasTPC, uint8_t); //! -DECLARE_SOA_COLUMN(HasTOF, hasTOF, uint8_t); //! -DECLARE_SOA_COLUMN(HasTRD, hasTRD, uint8_t); //! trdPattern() > 0 -DECLARE_SOA_COLUMN(HasEMCal, hasEMCal, uint8_t); //! trackEtaEmcal() in acceptance -DECLARE_SOA_COLUMN(HasHMPID, hasHMPID, uint8_t); //! matched in the sparse HMPID table - -// HMPID: sparse table, matched by hand per collision (see buildHmpidMap() in -// pidFeatureExtractor.cxx). -DECLARE_SOA_COLUMN(HmpidSignal, hmpidSignal, float); //! Cherenkov angle (rad), NaN if !hasHMPID -DECLARE_SOA_COLUMN(HmpidQMip, hmpidQMip, float); //! -DECLARE_SOA_COLUMN(HmpidNPhotons, hmpidNPhotons, int); //! -DECLARE_SOA_COLUMN(HmpidClusSize, hmpidClusSize, int); //! -DECLARE_SOA_COLUMN(HmpidMom, hmpidMom, float); //! - -// Bayesian PID posteriors. NaN when computeBayesianPid is false, or when the -// track doesn't have TPC (see computeBayesianProbs() in -// pidFeatureExtractor.cxx) - never a value that could be mistaken for a real -// posterior. -DECLARE_SOA_COLUMN(BayesProbPi, bayesProbPi, float); //! -DECLARE_SOA_COLUMN(BayesProbKa, bayesProbKa, float); //! -DECLARE_SOA_COLUMN(BayesProbPr, bayesProbPr, float); //! -DECLARE_SOA_COLUMN(BayesProbEl, bayesProbEl, float); //! - -// MC truth. -DECLARE_SOA_COLUMN(IsPhysicalPrimary, isPhysicalPrimary, uint8_t); //! -} // namespace pidfeat - -// Real/raw data: reconstructed features only. -DECLARE_SOA_TABLE(PidFeaturesData, "AOD", "PIDFEATD", //! - o2::soa::Index<>, - pidfeat::P, aod::track::Pt, pidfeat::Px, pidfeat::Py, pidfeat::Pz, - pidfeat::Eta, pidfeat::Phi, pidfeat::Sign, aod::track::TrackType, - pidfeat::Vz, aod::cent::CentFT0C, - aod::track::DcaXY, aod::track::DcaZ, - pidfeat::HasTPC, aod::track::TPCSignal, - pidtpc::TPCNSigmaPi, pidtpc::TPCNSigmaKa, pidtpc::TPCNSigmaPr, pidtpc::TPCNSigmaEl, - aod::track::TPCNClsFound, aod::track::TPCChi2NCl, - pidfeat::HasTOF, pidfeat::TofMass, aod::pidtofbeta::Beta, - pidtof::TOFNSigmaPi, pidtof::TOFNSigmaKa, pidtof::TOFNSigmaPr, pidtof::TOFNSigmaEl, - pidfeat::HasTRD, aod::track::TRDSignal, aod::track::TRDChi2, aod::track::TRDPattern, - aod::track::ITSClusterSizes, aod::track::ITSChi2NCl, - pidfeat::HasEMCal, aod::track::TrackEtaEMCAL, aod::track::TrackPhiEMCAL, - pidfeat::HasHMPID, pidfeat::HmpidSignal, pidfeat::HmpidQMip, - pidfeat::HmpidNPhotons, pidfeat::HmpidClusSize, pidfeat::HmpidMom, - pidfeat::BayesProbPi, pidfeat::BayesProbKa, pidfeat::BayesProbPr, pidfeat::BayesProbEl); - -// MC: reconstructed features + truth. Same reconstructed columns as -// PidFeaturesData plus PdgCode/IsPhysicalPrimary - kept as a distinct table -// (rather than optional columns on one table) because O2 tables have a -// fixed schema. -DECLARE_SOA_TABLE(PidFeaturesMc, "AOD", "MCPIDFEA", //! - o2::soa::Index<>, - pidfeat::P, aod::track::Pt, pidfeat::Px, pidfeat::Py, pidfeat::Pz, - pidfeat::Eta, pidfeat::Phi, pidfeat::Sign, aod::track::TrackType, - pidfeat::Vz, aod::cent::CentFT0C, - aod::track::DcaXY, aod::track::DcaZ, - pidfeat::HasTPC, aod::track::TPCSignal, - pidtpc::TPCNSigmaPi, pidtpc::TPCNSigmaKa, pidtpc::TPCNSigmaPr, pidtpc::TPCNSigmaEl, - aod::track::TPCNClsFound, aod::track::TPCChi2NCl, - pidfeat::HasTOF, pidfeat::TofMass, aod::pidtofbeta::Beta, - pidtof::TOFNSigmaPi, pidtof::TOFNSigmaKa, pidtof::TOFNSigmaPr, pidtof::TOFNSigmaEl, - pidfeat::HasTRD, aod::track::TRDSignal, aod::track::TRDChi2, aod::track::TRDPattern, - aod::track::ITSClusterSizes, aod::track::ITSChi2NCl, - pidfeat::HasEMCal, aod::track::TrackEtaEMCAL, aod::track::TrackPhiEMCAL, - pidfeat::HasHMPID, pidfeat::HmpidSignal, pidfeat::HmpidQMip, - pidfeat::HmpidNPhotons, pidfeat::HmpidClusSize, pidfeat::HmpidMom, - pidfeat::BayesProbPi, pidfeat::BayesProbKa, pidfeat::BayesProbPr, pidfeat::BayesProbEl, - aod::mcparticle::PdgCode, pidfeat::IsPhysicalPrimary); -} // namespace o2::aod - -#endif // PID_FEATURE_EXTRACTOR_H_ diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractorConfig.json b/Tools/PIDFeatureExtractor/pidFeatureExtractorConfig.json deleted file mode 100644 index f221462552a..00000000000 --- a/Tools/PIDFeatureExtractor/pidFeatureExtractorConfig.json +++ /dev/null @@ -1,150 +0,0 @@ -{ - "internal-dpl-clock": "", - - "internal-dpl-aod-reader": { - "time-limit": 0, - "aod-file-private": "AO2D.root", - "orbit-offset-enumeration": 0, - "orbit-multiplier-enumeration": 0, - "start-value-enumeration": 0, - "end-value-enumeration": -1, - "step-value-enumeration": 1 - }, - - "timestamp-task": { - "verbose": 0, - "rct-path": "RCT/Info/RunInformation", - "orbit-reset-path": "CTP/Calib/OrbitReset", - "ccdb-url": "http://alice-ccdb.cern.ch", - "isRun2MC": 0 - }, - - "bc-selection-task": { - "processRun2": 0, - "processRun3": 1 - }, - - "event-selection-task": { - "syst": "PbPb", - "muonSelection": 0, - "customDeltaBC": 0, - "isMC": 0, - "processRun2": 0, - "processRun3": 1 - }, - - "track-propagation": { - "ccdb-url": "http://alice-ccdb.cern.ch", - "grp-path": "GLO/GRP/GRP", - "grp-mag-path": "GLO/Config/GRPMagField", - "mVtxPath": "GLO/Calib/MeanVertex", - "geo-path": "GLO/Config/GeometryAligned", - "useMatLUT": 0, - "processStandard": 1, - "processCovariance": 0, - "processCovarianceMc": 0, - "minPropagationDistance": 83.1 - }, - - "track-selection-task": { - "isRun3": 1 - }, - - "pid-tpc-base": { - "ccdb-url": "http://alice-ccdb.cern.ch", - "parametrization-path": "TPC/Calib/Response", - "parametrization-el-path": "TPC/Calib/ResponseElectron", - "resoPath": "TPC/Calib/PIDResponse", - "ccdb-timestamp": 0, - "useNetworkCorrection": 0, - "autofetch-network": 1, - "enableNetworkOptimization": 1, - "networkPathLocally": "", - "networkPathCCDB": "Analysis/PID/TPC", - "onnxFile": "network.onnx", - "enableNetworkInference": 0 - }, - - "pid-tpc": { - "param-file": "", - "param-sigma": "TPC.PIDResponse.sigma:", - "ccdb-url": "http://alice-ccdb.cern.ch", - "ccdbPath": "TPC/Calib/PIDResponse" - }, - - "pid-tof-base": { - "ccdb-url": "http://alice-ccdb.cern.ch", - "parametrizationPath": "TOF/Calib/Response", - "passName": "", - "timeShiftCCDBPath": "", - "fatalOnPassNotAvailable": 1 - }, - - "pid-tof": { - "param-file": "", - "param-sigma": "TOF.PIDResponse.sigma:", - "ccdb-url": "http://alice-ccdb.cern.ch", - "ccdbPath": "TOF/Calib/Response", - "passName": "", - "timeShiftCCDBPath": "", - "parametrizationPath": "TOF/Calib/Response", - "fatalOnPassNotAvailable": 1 - }, - - "pid-tof-beta": { - "ccdb-url": "http://alice-ccdb.cern.ch" - }, - - "multiplicity-table": { - "doVertexZeq": 1, - "fractionOfEvents": 2, - "processRun2": 0, - "processRun3": 1 - }, - - "centrality-table": { - "ccdb-url": "http://alice-ccdb.cern.ch", - "ccdbPath": "Centrality", - "genName": "", - "processRun2": 0, - "processRun3": 1, - "doNotCrashOnNull": 1, - "processFV0A": 0, - "processFT0M": 0, - "processFT0A": 0, - "processFT0C": 1, - "processFDDM": 0, - "processNTPV": 0, - "processNGlobal": 0, - "processMFT": 0 - }, - - "pid-feature-extractor": { - "processData": 1, - "processMc": 0, - "etaMin": -99.0, - "etaMax": 99.0, - "ptMin": 0.0, - "ptMax": 9999.0, - "dcaxyMax": 9999.0, - "dcazMax": 9999.0, - "tpcMinClusters": 0, - "itsMinClusters": 4, - "computeBayesianPid": true, - "bayesianPriors": [1.0, 1.0, 1.0, 1.0], - "exportCsv": false, - "csvOutputPath": "pid_features" - }, - - "pid-onnx-inference": { - "processData": 1, - "processMc": 0, - "loadModelFromCcdb": true, - "ccdbUrl": "http://alice-ccdb.cern.ch", - "modelPathsCcdb": ["Users/YOURNAME/PidFeatureExtractor/model"], - "timestampCcdb": -1, - "onnxFileNames": ["pid_feature_model.onnx"], - "binsPtMl": [-1.0, 9999.0], - "nClassesMl": 4 - } -} diff --git a/Tools/PIDFeatureExtractor/run.sh b/Tools/PIDFeatureExtractor/run.sh deleted file mode 100644 index 41bd0e16243..00000000000 --- a/Tools/PIDFeatureExtractor/run.sh +++ /dev/null @@ -1,77 +0,0 @@ -#!/bin/bash - -# PID Feature Extractor Workflow (simplified, table-based) -# Detectors: TPC, TOF, TRD, ITS, EMCal, HMPID + centrality (FT0C) - the 7 -# detector groups / 34-feature contract used by the FSE PID model. -# Event level: Centrality FT0C (Pb-Pb Run 3) -# -# Output is two AOD-joinable tables (PidFeaturesData / PidFeaturesMc) -# written via the framework's own Produces<> mechanism; set the usual -# --aod-writer-json if you want to control where the output -# AnalysisResults-style file goes. CSV export alongside the table is -# available via exportCsv (off by default) - see pidFeatureExtractorConfig.json. -# -# DPG cuts (eta/pT/DCA/TPC-cluster) and Bayesian PID (TPC alone is enough; -# TOF folded in if present; priors configurable) are both optional, set in -# the same config block - see README.md for details. -# -# Mode is a JSON choice: pidFeatureExtractorConfig.json -> "pid-feature-extractor" -# "processData": 1, "processMc": 0 -> real/raw data (no MC truth) -# "processData": 0, "processMc": 1 -> MC (reconstructed + truth) -# This task does not itself abort if you leave both/neither on - the -# framework will just run whichever you've enabled or do nothing productive -# if neither is. Double check your config. - -CONFIG="$(pwd)/pidFeatureExtractorConfig.json" -OPTION="-b --configuration json://${CONFIG}" - -# CRITICAL: Add shared memory flag (~half your available system RAM) -SHM_SIZE="--shm-segment-size 4000000000" - -EXTRACTOR=~/alice/sw/BUILD/O2Physics-latest/O2Physics/stage/bin/o2-analysis-pid-feature-extractor -INFERENCE=~/alice/sw/BUILD/O2Physics-latest/O2Physics/stage/bin/o2-analysis-pid-onnx-inference - -echo "Starting O2Physics PID Feature Extraction + Inference Workflow..." -echo "Using configuration: ${CONFIG}" -echo "Shared memory segment size: ${SHM_SIZE}" -echo "" - -# Pipeline: -# timestamp → event selection → track propagation → track selection -# (needed for requireGlobalTrackInFilter()) -# → TPC PID → TOF PID → TOF beta -# → multiplicity → centrality (needed for CentFT0Cs, both MC and data) -# → feature extractor → ONNX inference -# -# TRD, ITS, EMCal: already in TracksExtra — no extra task needed -# HMPID: already in AO2D O2hmpid_001 — no extra task needed -# -# The inference task consumes PidFeaturesData/PidFeaturesMc directly (DPL -# wires the table dependency automatically since both tasks run in one -# workflow here) - drop the last pipe stage if you only want the features -# table and don't want to run inference. - -o2-analysis-timestamp ${OPTION} ${SHM_SIZE} | \ -o2-analysis-event-selection ${OPTION} ${SHM_SIZE} | \ -o2-analysis-track-propagation ${OPTION} ${SHM_SIZE} | \ -o2-analysis-trackselection ${OPTION} ${SHM_SIZE} | \ -o2-analysis-pid-tpc-base ${OPTION} ${SHM_SIZE} | \ -o2-analysis-pid-tpc ${OPTION} ${SHM_SIZE} | \ -o2-analysis-pid-tof-base ${OPTION} ${SHM_SIZE} | \ -o2-analysis-pid-tof ${OPTION} ${SHM_SIZE} | \ -o2-analysis-pid-tof-beta ${OPTION} ${SHM_SIZE} | \ -o2-analysis-multiplicity-table ${OPTION} ${SHM_SIZE} | \ -o2-analysis-centrality-table ${OPTION} ${SHM_SIZE} | \ -${EXTRACTOR} ${OPTION} ${SHM_SIZE} | \ -${INFERENCE} ${OPTION} ${SHM_SIZE} - -EXIT_CODE=$? - -echo "" -if [ $EXIT_CODE -eq 0 ]; then - echo "✓ Workflow completed successfully!" -else - echo "✗ Workflow failed with exit code: $EXIT_CODE" -fi - -exit $EXIT_CODE From d7ba0216386203b94fde97d340ba076135f46647 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Tue, 11 Aug 2026 21:39:11 +0100 Subject: [PATCH 13/16] Fix CI: clang-format, copyright header, magic numbers, configurable naming; rework pidOnnxInference as a proper adaptAnalysisTask --- Tools/PIDFeatureExtractor/CMakeLists.txt | 11 ++ Tools/PIDFeatureExtractor/README.md | 36 ++-- .../pidFeatureExtractor.cxx | 19 +- .../PIDFeatureExtractor/pidOnnxInference.cxx | 184 +++++++++--------- 4 files changed, 134 insertions(+), 116 deletions(-) diff --git a/Tools/PIDFeatureExtractor/CMakeLists.txt b/Tools/PIDFeatureExtractor/CMakeLists.txt index f5520cea85d..b163d0e9293 100644 --- a/Tools/PIDFeatureExtractor/CMakeLists.txt +++ b/Tools/PIDFeatureExtractor/CMakeLists.txt @@ -1,3 +1,14 @@ +# 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 diff --git a/Tools/PIDFeatureExtractor/README.md b/Tools/PIDFeatureExtractor/README.md index a1f87849e76..7afd9a539de 100644 --- a/Tools/PIDFeatureExtractor/README.md +++ b/Tools/PIDFeatureExtractor/README.md @@ -12,8 +12,10 @@ HMPID, plus event centrality). Two tasks: 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 - they're two -separate steps, not one pipeline (see "Running" below for why). +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 @@ -40,7 +42,7 @@ Mode is a runtime switch - enable `processData` for real data or | `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 | +| `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 | @@ -52,9 +54,15 @@ quality selection applied here rather than downstream. ## PidOnnxInference Takes the file `PidFeatureExtractor` wrote and runs the trained ONNX model -over it, row by row. The model can be loaded either from CCDB or from a -local file, which is handled by `o2::analysis::MlResponse` -(`Tools/ML/MlResponse.h`). +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 @@ -90,10 +98,10 @@ as an index: `0`=pion, `1`=kaon, `2`=proton, `3`=electron). ## Running -Both use the usual `--configuration json://your-config.json` mechanism. - -`PidFeatureExtractor` needs to run as part of the normal AOD pipeline, -since it reads track and collision data directly: +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 @@ -114,13 +122,15 @@ o2-analysis-timestamp --configuration json://$config_file -b | o2-analysis-pid-feature-extractor --configuration json://$config_file -b ``` -`PidOnnxInference` runs on its own, after that has finished - it just -opens the file the extractor wrote, so there's no AOD pipeline to build: +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-pid-onnx-inference --configuration json://$config_file -b +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 index 78b61e3c681..2e3f7b3ebc3 100644 --- a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx +++ b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx @@ -57,6 +57,11 @@ 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 @@ -85,8 +90,8 @@ int getItsNClusters(T const& track) { auto v = static_cast(track.itsClusterSizes()); int n = 0; - for (int layer = 0; layer < 7; layer++) { - if ((v >> (layer * 4)) & 0xF) { + for (int layer = 0; layer < kNumItsLayers; layer++) { + if ((v >> (layer * kBitsPerItsLayer)) & kItsLayerMask) { n++; } } @@ -154,8 +159,8 @@ struct PidFeatureExtractor { 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 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)"}; @@ -299,7 +304,7 @@ struct PidFeatureExtractor { } auto const& priors = bayesianPriors.value; float sum = 0.f; - for (int i = 0; i < 4; i++) { + for (int i = 0; i < kNumSpecies; i++) { float logL = -0.5f * nsTPC[i] * nsTPC[i]; if (hasTofIn) { logL += -0.5f * nsTOF[i] * nsTOF[i]; @@ -307,7 +312,7 @@ struct PidFeatureExtractor { out[i] = std::exp(logL) * priors[i]; sum += out[i]; } - for (int i = 0; i < 4; i++) { + for (int i = 0; i < kNumSpecies; i++) { out[i] = sum > 0.f ? out[i] / sum : 0.25f; } } @@ -370,7 +375,7 @@ struct PidFeatureExtractor { itsClusterSizes = track.itsClusterSizes(); itsChi2NCl = track.itsChi2NCl(); - hasEmcal = track.trackEtaEmcal() > -900.f; + hasEmcal = track.trackEtaEmcal() > kEmcalEtaOutOfAcceptance; trackEtaEmcal = track.trackEtaEmcal(); trackPhiEmcal = track.trackPhiEmcal(); diff --git a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx index 5f3956dca39..f3a10ad8e29 100644 --- a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx +++ b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx @@ -15,17 +15,17 @@ /// and write per-track class probabilities to a new ROOT file /// (and/or CSV). /// -/// Built as a plain DataProcessorSpec/AlgorithmSpec, NOT -/// adaptAnalysisTask<> - this is a standalone batch job with no AOD -/// table subscription, and adaptAnalysisTask's process() reflection -/// specifically requires an AOD table/iterator argument (confirmed -/// by a real compiler error: a raw ProcessingContext& process() -/// signature is rejected outright). Using the lower-level DPL -/// primitives directly avoids that machinery entirely, rather than -/// trying to satisfy a reflection mechanism built for a different -/// kind of task. Configurable<> member auto-binding is an -/// AnalysisTask-specific convenience, so options here are declared -/// explicitly and read via ic.options().get(...) instead. +/// 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 @@ -35,9 +35,8 @@ #include "Tools/ML/MlResponse.h" -#include -#include -#include +#include +#include #include #include @@ -58,14 +57,17 @@ 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 < 7; layer++) { - if ((v >> (layer * 4)) & 0xF) { + for (int layer = 0; layer < kNumItsLayers; layer++) { + if ((v >> (layer * kBitsPerItsLayer)) & kItsLayerMask) { n++; } } @@ -212,8 +214,6 @@ void runInference(std::string const& inputRootFile, std::string const& inputTree tree->GetEntry(i); // Effective presence = what the data says AND the group is enabled. - // Disabling a group here forces the same "absent" state as a real - // detector miss, regardless of what hasXXX says in the input tree. bool effTPC = hasTPC && groups.useTPC; bool effTOF = hasTOF && groups.useTOF; bool effTRD = hasTRD && groups.useTRD; @@ -294,83 +294,75 @@ void runInference(std::string const& inputRootFile, std::string const& inputTree } } // namespace -WorkflowSpec defineDataProcessing(ConfigContext const&) +/// 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", 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) { - DataProcessorSpec spec{ - "pid-onnx-inference", - Inputs{}, - Outputs{}, - AlgorithmSpec{[](InitContext& ic) { - auto inputRootFile = ic.options().get("inputRootFile"); - auto inputTreeName = ic.options().get("inputTreeName"); - auto outputPath = ic.options().get("outputPath"); - auto exportCsv = ic.options().get("exportCsv"); - auto loadModelFromCcdb = ic.options().get("loadModelFromCcdb"); - auto ccdbUrl = ic.options().get("ccdbUrl"); - auto modelPathsCcdb = ic.options().get>("modelPathsCcdb"); - auto timestampCcdb = ic.options().get("timestampCcdb"); - auto onnxFileNames = ic.options().get>("onnxFileNames"); - auto binsPtMl = ic.options().get>("binsPtMl"); - auto nClassesMl = static_cast(ic.options().get("nClassesMl")); - - GroupToggles groups; - groups.useTPC = ic.options().get("useTPC"); - groups.useTOF = ic.options().get("useTOF"); - groups.useTRD = ic.options().get("useTRD"); - groups.useITS = ic.options().get("useITS"); - groups.useEMCal = ic.options().get("useEMCal"); - groups.useHMPID = ic.options().get("useHMPID"); - groups.useCentrality = ic.options().get("useCentrality"); - - // 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; hardcoded here - // rather than exposed as options. - 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}; - - auto mlResponse = std::make_shared>(); - mlResponse->configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); - if (loadModelFromCcdb) { - auto ccdbApi = std::make_shared(); - ccdbApi->init(ccdbUrl); - mlResponse->setModelPathsCCDB(onnxFileNames, *ccdbApi, modelPathsCcdb, timestampCcdb); - } else { - mlResponse->setModelPathsLocal(onnxFileNames); - } - mlResponse->init(); - - runInference(inputRootFile, inputTreeName, outputPath, exportCsv, *mlResponse, groups); - - return [](ProcessingContext& pc) { - // One-shot batch job: all work already happened in init(). Signal - // completion immediately rather than waiting for input that will - // never arrive. - pc.services().get().endOfStream(); - pc.services().get().readyToQuit(QuitRequest::Me); - }; - }}, - Options{ - {"inputRootFile", VariantType::String, "pid_features_data.root", {"ROOT file produced by pidFeatureExtractor.cxx"}}, - {"inputTreeName", VariantType::String, "pid_features", {"Name of the TTree inside inputRootFile"}}, - {"outputPath", VariantType::String, "pid_predictions", {"Output file base name (no extension)"}}, - {"exportCsv", VariantType::Bool, false, {"Also write predictions to CSV alongside the ROOT output"}}, - {"loadModelFromCcdb", VariantType::Bool, true, {"Load the ONNX model from CCDB (else from onnxFileNames as a local path)"}}, - {"ccdbUrl", VariantType::String, "http://alice-ccdb.cern.ch", {"CCDB URL"}}, - {"modelPathsCcdb", VariantType::ArrayString, std::vector{"Users/YOURNAME/PidFeatureExtractor/model"}, {"CCDB path to the model"}}, - {"timestampCcdb", VariantType::Int64, static_cast(-1), {"CCDB query timestamp for the model, -1 = latest"}}, - {"onnxFileNames", VariantType::ArrayString, std::vector{"pid_feature_model.onnx"}, {"Local ONNX file path(s), used when loadModelFromCcdb is false"}}, - {"binsPtMl", VariantType::ArrayDouble, std::vector{-1., 9999.}, {"pT bin edges for MlResponse (single bin = model isn't pT-binned)"}}, - {"nClassesMl", VariantType::Int, kNumClasses, {"Number of model output classes"}}, - {"useTPC", VariantType::Bool, true, {"Include TPC in inference. Default true (all detectors present); set false to force TPC excluded regardless of the data"}}, - {"useTOF", VariantType::Bool, true, {"Include TOF in inference"}}, - {"useTRD", VariantType::Bool, true, {"Include TRD in inference"}}, - {"useITS", VariantType::Bool, true, {"Include ITS in inference"}}, - {"useEMCal", VariantType::Bool, true, {"Include EMCal in inference"}}, - {"useHMPID", VariantType::Bool, true, {"Include HMPID in inference"}}, - {"useCentrality", VariantType::Bool, true, {"Include centrality in inference"}}, - }}; - - return WorkflowSpec{spec}; + return WorkflowSpec{adaptAnalysisTask(cfgc)}; } From 9459177aacced3d22638daa01ecdc3bfd4c20468 Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Tue, 11 Aug 2026 21:48:57 +0100 Subject: [PATCH 14/16] Fix clang-format indentation, cppcheck const-array warning, duplicate README heading --- Tools/PIDFeatureExtractor/README.md | 4 ++-- Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx | 2 +- Tools/PIDFeatureExtractor/pidOnnxInference.cxx | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/Tools/PIDFeatureExtractor/README.md b/Tools/PIDFeatureExtractor/README.md index 7afd9a539de..2ad58a3fc4b 100644 --- a/Tools/PIDFeatureExtractor/README.md +++ b/Tools/PIDFeatureExtractor/README.md @@ -33,7 +33,7 @@ cuts, writes one row containing: Mode is a runtime switch - enable `processData` for real data or `processMc` for MC (reconstructed + truth), not both. -### Configurable options +### PidFeatureExtractor options | Option | Default | What it does | |---|---|---| @@ -71,7 +71,7 @@ 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. -### Configurable options +### PidOnnxInference options | Option | Default | What it does | |---|---|---| diff --git a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx index 2e3f7b3ebc3..b7c58ebe9f8 100644 --- a/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx +++ b/Tools/PIDFeatureExtractor/pidFeatureExtractor.cxx @@ -296,7 +296,7 @@ struct PidFeatureExtractor { /// 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, float nsTPC[4], bool hasTofIn, float nsTOF[4], float out[4]) const + 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; diff --git a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx index f3a10ad8e29..ac38c30a426 100644 --- a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx +++ b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx @@ -115,8 +115,8 @@ struct GroupToggles { /// (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::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()) { From 186bf3c2ec71c19d41f57fbdb10864969bf7ef21 Mon Sep 17 00:00:00 2001 From: ALICE Action Bot Date: Tue, 11 Aug 2026 20:54:04 +0000 Subject: [PATCH 15/16] MegaLinter fixes --- Tools/PIDFeatureExtractor/README.md | 60 ++++++++++++++--------------- 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/Tools/PIDFeatureExtractor/README.md b/Tools/PIDFeatureExtractor/README.md index 2ad58a3fc4b..c0773c63fd4 100644 --- a/Tools/PIDFeatureExtractor/README.md +++ b/Tools/PIDFeatureExtractor/README.md @@ -35,18 +35,18 @@ Mode is a runtime switch - enable `processData` for real data or ### 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 | +| 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. @@ -73,24 +73,24 @@ 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 | +| 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, From d056ad4f01a5b4aaeabfffe1023953e0cc90328f Mon Sep 17 00:00:00 2001 From: Robert Forynski Date: Wed, 12 Aug 2026 14:45:49 +0100 Subject: [PATCH 16/16] Fix Configurable const-lvalue binding error --- Tools/PIDFeatureExtractor/pidOnnxInference.cxx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx index ac38c30a426..86030e3ff82 100644 --- a/Tools/PIDFeatureExtractor/pidOnnxInference.cxx +++ b/Tools/PIDFeatureExtractor/pidOnnxInference.cxx @@ -314,7 +314,7 @@ struct PidOnnxInference { 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", kNumClasses, "Number of model output classes"}; + 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"};