Skip to content

Commit 86aabca

Browse files
forynskiforynski
committed
Add per-detector-group toggles to inference; rewrite README to match Tools/PIDML style
1 parent f21f33e commit 86aabca

2 files changed

Lines changed: 215 additions & 42 deletions

File tree

‎Tools/PIDFeatureExtractor/README.md‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
# PID Feature Extractor + ONNX Inference
2+
3+
This provides particle identification for ALICE Run 3 Pb-Pb analyses using
4+
a trained ML model (a detector-aware attention model conditioned on which
5+
detectors each track actually has hits in - TPC, TOF, TRD, ITS, EMCal,
6+
HMPID, plus event centrality). Two tasks:
7+
8+
- **`pidFeatureExtractor.cxx`** reads AO2D data and writes out the model's
9+
input features - kinematics, per-detector PID signals, and detector
10+
presence flags - to a ROOT file (and optionally CSV).
11+
- **`pidOnnxInference.cxx`** takes that file, runs the trained ONNX model
12+
over it, and writes back a probability for each particle species
13+
(pion / kaon / proton / electron) per track.
14+
15+
You run the extractor first, then inference on its output - they're two
16+
separate steps, not one pipeline (see "Running" below for why).
17+
18+
## PidFeatureExtractor
19+
20+
An ordinary AOD-subscribing analysis task. It reads track and collision
21+
data and, for each track passing the (optional, off by default) quality
22+
cuts, writes one row containing:
23+
24+
- kinematics (momentum, eta, phi, DCA)
25+
- per-detector signals for TPC, TOF, TRD, ITS, EMCal, and HMPID, each with
26+
a flag saying whether that detector actually has a hit on this track
27+
- event centrality
28+
- a Bayesian PID posterior, for comparison against the ML model
29+
- for MC only: the true particle ID and whether it's a physical primary
30+
31+
Mode is a runtime switch - enable `processData` for real data or
32+
`processMc` for MC (reconstructed + truth), not both.
33+
34+
### Configurable options
35+
36+
| Option | Default | What it does |
37+
|---|---|---|
38+
| `outputPath` | `pid_features` | Output file base name |
39+
| `exportROOT` | `true` | Write a ROOT file |
40+
| `exportCsv` | `false` | Also write CSV |
41+
| `etaMin` / `etaMax` | `-99` / `99` | Eta cut - wide open by default (no cut) |
42+
| `ptMin` / `ptMax` | `0` / `9999` | pT cut, GeV/c - wide open by default |
43+
| `dcaxyMax` / `dcazMax` | `9999` / `9999` | DCA cuts, cm - wide open by default |
44+
| `itsMinClusters` | `0` | Minimum ITS clusters - `0` = no cut |
45+
| `tpcMinClusters` | `0` | Minimum TPC clusters - `0` = no cut |
46+
| `computeBayesianPid` | `true` | Compute the comparison Bayesian posterior |
47+
| `bayesianPriors` | flat (`1,1,1,1`) | Per-species priors `[pi, ka, pr, el]` for the Bayesian posterior |
48+
49+
All the cuts default to "off" - tighten them in your config if you want
50+
quality selection applied here rather than downstream.
51+
52+
## PidOnnxInference
53+
54+
Takes the file `PidFeatureExtractor` wrote and runs the trained ONNX model
55+
over it, row by row. The model can be loaded either from CCDB or from a
56+
local file, which is handled by `o2::analysis::MlResponse`
57+
(`Tools/ML/MlResponse.h`).
58+
59+
By default it assumes every detector group is present and usable, exactly
60+
as the input data says. If you want to see how the model behaves with a
61+
detector deliberately left out - for testing, or to match a specific
62+
detector configuration - each group can be switched off independently;
63+
turning one off overrides the data for that group, the same way a genuine
64+
detector miss would look.
65+
66+
### Configurable options
67+
68+
| Option | Default | What it does |
69+
|---|---|---|
70+
| `inputRootFile` | `pid_features_data.root` | File written by `PidFeatureExtractor` |
71+
| `inputTreeName` | `pid_features` | Tree name inside it |
72+
| `outputPath` | `pid_predictions` | Output file base name |
73+
| `exportCsv` | `false` | Also write CSV |
74+
| `loadModelFromCcdb` | `true` | Load the model from CCDB; set `false` to use a local file instead |
75+
| `ccdbUrl` | `http://alice-ccdb.cern.ch` | |
76+
| `modelPathsCcdb` | *(placeholder)* | CCDB path to your model - set this to a real path before running |
77+
| `timestampCcdb` | `-1` | `-1` = latest |
78+
| `onnxFileNames` | `pid_feature_model.onnx` | Local model file, used when `loadModelFromCcdb` is `false` |
79+
| `useTPC` | `true` | Include TPC. Set `false` to exclude it from inference regardless of the data |
80+
| `useTOF` | `true` | Include TOF |
81+
| `useTRD` | `true` | Include TRD |
82+
| `useITS` | `true` | Include ITS |
83+
| `useEMCal` | `true` | Include EMCal |
84+
| `useHMPID` | `true` | Include HMPID |
85+
| `useCentrality` | `true` | Include event centrality |
86+
87+
Output columns are `mlProbPi`, `mlProbKa`, `mlProbPr`, `mlProbEl` (one
88+
probability per species) and `mlPredictedClass` (the most likely species,
89+
as an index: `0`=pion, `1`=kaon, `2`=proton, `3`=electron).
90+
91+
## Running
92+
93+
Both use the usual `--configuration json://your-config.json` mechanism.
94+
95+
`PidFeatureExtractor` needs to run as part of the normal AOD pipeline,
96+
since it reads track and collision data directly:
97+
98+
```bash
99+
#!/bin/bash
100+
101+
config_file="my-config.json"
102+
103+
o2-analysis-timestamp --configuration json://$config_file -b |
104+
o2-analysis-event-selection --configuration json://$config_file -b |
105+
o2-analysis-track-propagation --configuration json://$config_file -b |
106+
o2-analysis-trackselection --configuration json://$config_file -b |
107+
o2-analysis-pid-tpc-base --configuration json://$config_file -b |
108+
o2-analysis-pid-tpc --configuration json://$config_file -b |
109+
o2-analysis-pid-tof-base --configuration json://$config_file -b |
110+
o2-analysis-pid-tof --configuration json://$config_file -b |
111+
o2-analysis-pid-tof-beta --configuration json://$config_file -b |
112+
o2-analysis-multiplicity-table --configuration json://$config_file -b |
113+
o2-analysis-centrality-table --configuration json://$config_file -b |
114+
o2-analysis-pid-feature-extractor --configuration json://$config_file -b
115+
```
116+
117+
`PidOnnxInference` runs on its own, after that has finished - it just
118+
opens the file the extractor wrote, so there's no AOD pipeline to build:
119+
120+
```bash
121+
#!/bin/bash
122+
123+
config_file="my-config.json"
124+
125+
o2-analysis-pid-onnx-inference --configuration json://$config_file -b
126+
```

‎Tools/PIDFeatureExtractor/pidOnnxInference.cxx‎

Lines changed: 89 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545

4646
#include <cstdint>
4747
#include <fstream>
48+
#include <limits>
4849
#include <memory>
4950
#include <string>
5051
#include <vector>
@@ -56,6 +57,7 @@ using namespace o2::framework;
5657
namespace
5758
{
5859
constexpr int kNumClasses = 4; // pi, ka, pr, el - fixed order throughout, matches the paper's model
60+
constexpr float kNaN = std::numeric_limits<float>::quiet_NaN();
5961

6062
/// itsClusterSizes packs 7 ITS layers into 4 bits each; a derived cluster
6163
/// count is a far more sensible model input than the raw packed value.
@@ -81,18 +83,38 @@ int argmax4(std::vector<float> const& v)
8183
return best;
8284
}
8385

86+
/// Per-group enable/disable, independent of what the input tree's hasXXX
87+
/// flags say. Default is everything enabled (true) - the normal case,
88+
/// using each track's real detector coverage as-is. Turning a group off
89+
/// forces its features to the same "absent" sentinel used when the
90+
/// detector genuinely didn't fire, and clears its mask bit - useful for
91+
/// testing how the model behaves with a detector deliberately excluded,
92+
/// independent of the data itself.
93+
struct GroupToggles {
94+
bool useTPC = true;
95+
bool useTOF = true;
96+
bool useTRD = true;
97+
bool useITS = true;
98+
bool useEMCal = true;
99+
bool useHMPID = true;
100+
bool useCentrality = true;
101+
};
102+
84103
/// All the real work: load the model, read the whole input tree, run
85104
/// inference row by row, write predictions. Called once from init().
86105
///
87106
/// Feature order fed to the model - THIS MUST MATCH YOUR TRAINING SCRIPT'S
88107
/// COLUMN ORDER EXACTLY. Reasonable default (every reconstructed feature
89108
/// except vz/centFT0C/sign/trackType and the Bayesian columns, which are a
90109
/// comparison baseline, not a model input), followed by a 7-length group
91-
/// mask (TPC/TOF/TRD/ITS/EMCal/HMPID/centrality; ITS and centrality are
92-
/// assumed always-present). NOT verified against your actual training code.
110+
/// mask (TPC/TOF/TRD/ITS/EMCal/HMPID/centrality). Each group can be
111+
/// disabled via GroupToggles regardless of what the data says - see there
112+
/// for details. ITS and centrality have no hasXXX flag in the input tree
113+
/// (assumed always-present in the data itself), so their toggle is the
114+
/// only way to exclude them.
93115
void runInference(std::string const& inputRootFile, std::string const& inputTreeName,
94116
std::string const& outputPath, bool exportCsv,
95-
o2::analysis::MlResponse<float>& mlResponse)
117+
o2::analysis::MlResponse<float>& mlResponse, GroupToggles const& groups)
96118
{
97119
std::unique_ptr<TFile> inFile(TFile::Open(inputRootFile.c_str(), "READ"));
98120
if (!inFile || inFile->IsZombie()) {
@@ -189,8 +211,17 @@ void runInference(std::string const& inputRootFile, std::string const& inputTree
189211
for (Long64_t i = 0; i < nEntries; i++) {
190212
tree->GetEntry(i);
191213

214+
// Effective presence = what the data says AND the group is enabled.
215+
// Disabling a group here forces the same "absent" state as a real
216+
// detector miss, regardless of what hasXXX says in the input tree.
217+
bool effTPC = hasTPC && groups.useTPC;
218+
bool effTOF = hasTOF && groups.useTOF;
219+
bool effTRD = hasTRD && groups.useTRD;
220+
bool effEMCal = hasEMCal && groups.useEMCal;
221+
bool effHMPID = hasHMPID && groups.useHMPID;
222+
192223
x.clear();
193-
x.reserve(38 + 7);
224+
x.reserve(39 + 7);
194225
x.push_back(p);
195226
x.push_back(pt);
196227
x.push_back(px);
@@ -200,44 +231,44 @@ void runInference(std::string const& inputRootFile, std::string const& inputTree
200231
x.push_back(phi);
201232
x.push_back(dcaXY);
202233
x.push_back(dcaZ);
203-
x.push_back(static_cast<float>(hasTPC));
204-
x.push_back(tpcSignal);
205-
x.push_back(tpcNSigmaPi);
206-
x.push_back(tpcNSigmaKa);
207-
x.push_back(tpcNSigmaPr);
208-
x.push_back(tpcNSigmaEl);
209-
x.push_back(static_cast<float>(tpcNClsFound));
210-
x.push_back(tpcChi2NCl);
211-
x.push_back(static_cast<float>(hasTOF));
212-
x.push_back(tofMass);
213-
x.push_back(beta);
214-
x.push_back(tofNSigmaPi);
215-
x.push_back(tofNSigmaKa);
216-
x.push_back(tofNSigmaPr);
217-
x.push_back(tofNSigmaEl);
218-
x.push_back(static_cast<float>(hasTRD));
219-
x.push_back(trdSignal);
220-
x.push_back(trdChi2);
221-
x.push_back(static_cast<float>(trdPattern));
222-
x.push_back(static_cast<float>(getItsNClusters(static_cast<uint32_t>(itsClusterSizes))));
223-
x.push_back(itsChi2NCl);
224-
x.push_back(static_cast<float>(hasEMCal));
225-
x.push_back(trackEtaEmcal);
226-
x.push_back(trackPhiEmcal);
227-
x.push_back(static_cast<float>(hasHMPID));
228-
x.push_back(hmpidSignal);
229-
x.push_back(hmpidQMip);
230-
x.push_back(static_cast<float>(hmpidNPhotons));
231-
x.push_back(static_cast<float>(hmpidClusSize));
232-
x.push_back(hmpidMom);
234+
x.push_back(static_cast<float>(effTPC));
235+
x.push_back(effTPC ? tpcSignal : kNaN);
236+
x.push_back(effTPC ? tpcNSigmaPi : kNaN);
237+
x.push_back(effTPC ? tpcNSigmaKa : kNaN);
238+
x.push_back(effTPC ? tpcNSigmaPr : kNaN);
239+
x.push_back(effTPC ? tpcNSigmaEl : kNaN);
240+
x.push_back(effTPC ? static_cast<float>(tpcNClsFound) : 0.f);
241+
x.push_back(effTPC ? tpcChi2NCl : kNaN);
242+
x.push_back(static_cast<float>(effTOF));
243+
x.push_back(effTOF ? tofMass : kNaN);
244+
x.push_back(effTOF ? beta : kNaN);
245+
x.push_back(effTOF ? tofNSigmaPi : kNaN);
246+
x.push_back(effTOF ? tofNSigmaKa : kNaN);
247+
x.push_back(effTOF ? tofNSigmaPr : kNaN);
248+
x.push_back(effTOF ? tofNSigmaEl : kNaN);
249+
x.push_back(static_cast<float>(effTRD));
250+
x.push_back(effTRD ? trdSignal : kNaN);
251+
x.push_back(effTRD ? trdChi2 : kNaN);
252+
x.push_back(effTRD ? static_cast<float>(trdPattern) : 0.f);
253+
x.push_back(groups.useITS ? static_cast<float>(getItsNClusters(static_cast<uint32_t>(itsClusterSizes))) : 0.f);
254+
x.push_back(groups.useITS ? itsChi2NCl : kNaN);
255+
x.push_back(static_cast<float>(effEMCal));
256+
x.push_back(effEMCal ? trackEtaEmcal : kNaN);
257+
x.push_back(effEMCal ? trackPhiEmcal : kNaN);
258+
x.push_back(static_cast<float>(effHMPID));
259+
x.push_back(effHMPID ? hmpidSignal : kNaN);
260+
x.push_back(effHMPID ? hmpidQMip : kNaN);
261+
x.push_back(effHMPID ? static_cast<float>(hmpidNPhotons) : 0.f);
262+
x.push_back(effHMPID ? static_cast<float>(hmpidClusSize) : 0.f);
263+
x.push_back(effHMPID ? hmpidMom : kNaN);
233264
// 7-length group mask
234-
x.push_back(static_cast<float>(hasTPC));
235-
x.push_back(static_cast<float>(hasTOF));
236-
x.push_back(static_cast<float>(hasTRD));
237-
x.push_back(1.f); // ITS
238-
x.push_back(static_cast<float>(hasEMCal));
239-
x.push_back(static_cast<float>(hasHMPID));
240-
x.push_back(1.f); // centrality
265+
x.push_back(static_cast<float>(effTPC));
266+
x.push_back(static_cast<float>(effTOF));
267+
x.push_back(static_cast<float>(effTRD));
268+
x.push_back(static_cast<float>(groups.useITS));
269+
x.push_back(static_cast<float>(effEMCal));
270+
x.push_back(static_cast<float>(effHMPID));
271+
x.push_back(static_cast<float>(groups.useCentrality));
241272

242273
mlResponse.isSelectedMl(x, pt, mlOutput); // return value (selection) unused; mlOutput carries the 4 raw scores
243274
mlProbPi = mlOutput[0];
@@ -282,6 +313,15 @@ WorkflowSpec defineDataProcessing(ConfigContext const&)
282313
auto binsPtMl = ic.options().get<std::vector<double>>("binsPtMl");
283314
auto nClassesMl = static_cast<int8_t>(ic.options().get<int>("nClassesMl"));
284315

316+
GroupToggles groups;
317+
groups.useTPC = ic.options().get<bool>("useTPC");
318+
groups.useTOF = ic.options().get<bool>("useTOF");
319+
groups.useTRD = ic.options().get<bool>("useTRD");
320+
groups.useITS = ic.options().get<bool>("useITS");
321+
groups.useEMCal = ic.options().get<bool>("useEMCal");
322+
groups.useHMPID = ic.options().get<bool>("useHMPID");
323+
groups.useCentrality = ic.options().get<bool>("useCentrality");
324+
285325
// Unused thresholds (CutNot everywhere) - this task always reports
286326
// all four probabilities rather than applying a selection cut, so
287327
// cutsMl/cutDirMl don't need to be user-configurable; hardcoded here
@@ -301,7 +341,7 @@ WorkflowSpec defineDataProcessing(ConfigContext const&)
301341
}
302342
mlResponse->init();
303343

304-
runInference(inputRootFile, inputTreeName, outputPath, exportCsv, *mlResponse);
344+
runInference(inputRootFile, inputTreeName, outputPath, exportCsv, *mlResponse, groups);
305345

306346
return [](ProcessingContext& pc) {
307347
// One-shot batch job: all work already happened in init(). Signal
@@ -323,6 +363,13 @@ WorkflowSpec defineDataProcessing(ConfigContext const&)
323363
{"onnxFileNames", VariantType::ArrayString, std::vector<std::string>{"pid_feature_model.onnx"}, {"Local ONNX file path(s), used when loadModelFromCcdb is false"}},
324364
{"binsPtMl", VariantType::ArrayDouble, std::vector<double>{-1., 9999.}, {"pT bin edges for MlResponse (single bin = model isn't pT-binned)"}},
325365
{"nClassesMl", VariantType::Int, kNumClasses, {"Number of model output classes"}},
366+
{"useTPC", VariantType::Bool, true, {"Include TPC in inference. Default true (all detectors present); set false to force TPC excluded regardless of the data"}},
367+
{"useTOF", VariantType::Bool, true, {"Include TOF in inference"}},
368+
{"useTRD", VariantType::Bool, true, {"Include TRD in inference"}},
369+
{"useITS", VariantType::Bool, true, {"Include ITS in inference"}},
370+
{"useEMCal", VariantType::Bool, true, {"Include EMCal in inference"}},
371+
{"useHMPID", VariantType::Bool, true, {"Include HMPID in inference"}},
372+
{"useCentrality", VariantType::Bool, true, {"Include centrality in inference"}},
326373
}};
327374

328375
return WorkflowSpec{spec};

0 commit comments

Comments
 (0)