Skip to content

Commit df8948f

Browse files
forynskialibuild
forynski
and
alibuild
authored
[Tools] Add PID feature extractor and ONNX inference tasks (#17419)
Co-authored-by: ALICE Action Bot <alibuild@cern.ch>
1 parent 5763e07 commit df8948f

5 files changed

Lines changed: 1022 additions & 0 deletions

File tree

‎Tools/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,4 @@
1212
add_subdirectory(PIDML)
1313
add_subdirectory(ML)
1414
add_subdirectory(KFparticle)
15+
add_subdirectory(PIDFeatureExtractor)

‎Tools/PIDFeatureExtractor/CMakeLists.txt‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
# Copyright 2019-2020 CERN and copyright holders of ALICE O2.
2+
# See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3+
# All rights not expressly granted are reserved.
4+
#
5+
# This software is distributed under the terms of the GNU General Public
6+
# License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7+
#
8+
# In applying this license CERN does not waive the privileges and immunities
9+
# granted to it by virtue of its status as an Intergovernmental Organization
10+
# or submit itself to any jurisdiction.
11+
12+
o2physics_add_dpl_workflow(pid-feature-extractor
13+
SOURCES pidFeatureExtractor.cxx
14+
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
15+
COMPONENT_NAME Analysis)
16+
17+
o2physics_add_dpl_workflow(pid-onnx-inference
18+
SOURCES pidOnnxInference.cxx
19+
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore
20+
COMPONENT_NAME Analysis)

‎Tools/PIDFeatureExtractor/README.md‎

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

0 commit comments

Comments
 (0)