Skip to content

Commit e3eaeb5

Browse files
ChSonnabendChSonnabend
committed
Adjust secondary classifcation
1 parent 94259cf commit e3eaeb5

6 files changed

Lines changed: 219 additions & 53 deletions

File tree

‎DataFormats/simulation/include/SimulationDataFormat/StackParam.h‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,10 @@ struct StackParam : public o2::conf::ConfigurableParamHelper<StackParam> {
2929
std::string transportPrimaryFuncName = "";
3030
bool transportPrimaryInvert = false;
3131
// Used when transportPrimary="onnx". The model is fetched as raw ONNX bytes
32-
// and class 1 means "skip GEANT transport".
32+
// and class 1 means "this track and all descendants produce no hits".
33+
// Despite the legacy parameter name, ONNX runs at PreTrack for primaries
34+
// AND secondaries. Output is a single [batch, scores] float tensor; select
35+
// the class-1 score below. Disable ApplySigmoid for probability outputs.
3336
std::string transportPrimaryOnnxCCDBPath = "";
3437
float transportPrimaryOnnxThreshold = 0.5f;
3538
int transportPrimaryOnnxOutputIndex = 0;

‎Detectors/Base/include/DetectorsBase/Stack.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,10 @@ class Stack : public FairGenericStack
177177

178178
std::vector<MCTrack> const* const getMCTracks() const { return mTracks; }
179179

180+
/// Classify a birth track at PreTrack; false means stop transport before any hit.
181+
bool transportTrack(const TParticle& particle, double eventX, double eventY, double eventZ);
182+
bool hasTrackTransportModel() const { return static_cast<bool>(mTransportTrack); }
183+
180184
/// Clone for worker (used in MT mode only)
181185
FairGenericStack* CloneStack() const override;
182186

@@ -301,6 +305,8 @@ class Stack : public FairGenericStack
301305

302306
TransportFcn mTransportPrimary = [](const TParticle& p, const std::vector<TParticle>& particles) { return false; }; //! a function to inhibit the tracking of a particle
303307

308+
std::function<bool(const TParticle&, double, double, double)> mTransportTrack; //! ONNX decision for primaries and secondaries
309+
304310
// storage for track references
305311
std::vector<o2::TrackReference>* mTrackRefs = nullptr; //!
306312

‎Detectors/Base/include/DetectorsBase/TrackTransportUtils.h‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
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+
/// \file Stack.h
13+
/// \brief Definition of the Stack class
14+
/// \author M. Al-Turany - June 2014
15+
16+
#ifndef O2_TRACK_TRANSPORT_UTILS_H
17+
#define O2_TRACK_TRANSPORT_UTILS_H
18+
19+
#include "SimulationDataFormat/O2DatabasePDG.h"
20+
#include <TParticle.h>
21+
#include <TParticlePDG.h>
22+
#include <algorithm>
23+
#include <cmath>
24+
#include <limits>
25+
#include <stdexcept>
26+
#include <vector>
27+
28+
namespace o2::data::detail
29+
{
30+
// Same feature order, units, PDG masses and missing values as
31+
// extract_o2_kine_training.C. Normalisation/imputation belongs in the graph.
32+
inline std::vector<float> makeTrackTransportFeatures(const TParticle& particle, double eventX, double eventY, double eventZ)
33+
{
34+
const double px = particle.Px();
35+
const double py = particle.Py();
36+
const double pz = particle.Pz();
37+
const double momentum = std::sqrt(px * px + py * py + pz * pz);
38+
const double pt = std::hypot(px, py);
39+
bool massKnown = false;
40+
double mass = o2::O2DatabasePDG::Mass(particle.GetPdgCode(), massKnown);
41+
const auto* pdgInfo = particle.GetPDG();
42+
if (!massKnown) mass = pdgInfo ? pdgInfo->Mass() : 0.;
43+
const double missing = std::numeric_limits<double>::quiet_NaN();
44+
const double energy = std::sqrt(std::max(0., mass * mass + momentum * momentum));
45+
const double eta = momentum > std::abs(pz) ? 0.5 * std::log((momentum + pz) / (momentum - pz)) : missing;
46+
const double theta = momentum > 0. ? std::acos(pz / momentum) : missing;
47+
const double rapidity = energy > std::abs(pz) ? 0.5 * std::log((energy + pz) / (energy - pz)) : missing;
48+
double chargeSign = pdgInfo == nullptr || pdgInfo->Charge() == 0. ? 0. : std::copysign(1., pdgInfo->Charge());
49+
const int code = particle.GetPdgCode();
50+
if (std::abs(code) >= 1000000000) {
51+
chargeSign = (std::abs(code) / 10000) % 1000 == 0 ? 0. : (code > 0 ? 1. : -1.);
52+
}
53+
const double dx = particle.Vx() - eventX;
54+
const double dy = particle.Vy() - eventY;
55+
const double dz = particle.Vz() - eventZ;
56+
const double pdg = particle.GetPdgCode();
57+
58+
return {static_cast<float>(pdg), static_cast<float>(std::abs(pdg)), static_cast<float>(chargeSign),
59+
static_cast<float>(mass), static_cast<float>(energy), static_cast<float>(energy - mass),
60+
static_cast<float>(px), static_cast<float>(py), static_cast<float>(pz), static_cast<float>(momentum),
61+
static_cast<float>(pt), static_cast<float>(eta), static_cast<float>(std::atan2(py, px)), static_cast<float>(theta),
62+
static_cast<float>(rapidity), static_cast<float>(particle.Vx()), static_cast<float>(particle.Vy()),
63+
static_cast<float>(particle.Vz()), static_cast<float>(particle.T() * 1.e9), static_cast<float>(dx),
64+
static_cast<float>(dy), static_cast<float>(dz), static_cast<float>(std::hypot(particle.Vx(), particle.Vy())),
65+
static_cast<float>(std::hypot(dx, dy)), static_cast<float>(std::sqrt(dx * dx + dy * dy + dz * dz))};
66+
}
67+
68+
69+
inline bool transportFromOnnxScore(float score, float threshold, bool applySigmoid)
70+
{
71+
if (!std::isfinite(threshold) || threshold < 0.f || threshold > 1.f) {
72+
throw std::runtime_error("ONNX pruning threshold must be finite and in [0,1]");
73+
}
74+
if (!std::isfinite(score)) {
75+
throw std::runtime_error("ONNX pruning returned a non-finite score; refusing to reject a track");
76+
}
77+
if (applySigmoid) {
78+
score = score >= 0.f ? 1.f / (1.f + std::exp(-score)) : std::exp(score) / (1.f + std::exp(score));
79+
}
80+
if (score < 0.f || score > 1.f) {
81+
throw std::runtime_error("ONNX pruning probability is outside [0,1]");
82+
}
83+
// Class 1: this track and all descendants produce zero detector hits.
84+
return score < threshold;
85+
}
86+
} // namespace o2::data::detail
87+
#endif

‎Detectors/Base/src/Stack.cxx‎

Lines changed: 67 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
#include "CommonUtils/ConfigurationMacroHelper.h"
2828
#include "CCDB/BasicCCDBManager.h"
2929
#include "ML/OrtInterface.h"
30+
#include "DetectorsBase/TrackTransportUtils.h"
3031

3132
#include "TLorentzVector.h" // for TLorentzVector
3233
#include "TParticle.h" // for TParticle
@@ -41,6 +42,8 @@
4142
#include <cmath>
4243
#include <map>
4344
#include <memory>
45+
#include <limits>
46+
#include <mutex>
4447
#include <stdexcept>
4548
#include <unordered_map>
4649

@@ -58,14 +61,17 @@ namespace
5861
// ONNX graph, keeping this code independent of model topology.
5962
constexpr size_t OnnxFeatureCount = 25;
6063

61-
class OnnxPrimaryTransport
64+
class OnnxTrackTransport
6265
{
6366
public:
64-
explicit OnnxPrimaryTransport(const o2::sim::StackParam& param)
67+
explicit OnnxTrackTransport(const o2::sim::StackParam& param)
6568
: mThreshold(param.transportPrimaryOnnxThreshold),
6669
mOutputIndex(param.transportPrimaryOnnxOutputIndex),
6770
mApplySigmoid(param.transportPrimaryOnnxApplySigmoid)
6871
{
72+
if (!std::isfinite(mThreshold) || mThreshold < 0.f || mThreshold > 1.f) {
73+
throw std::runtime_error("ONNX pruning threshold must be finite and in [0,1]");
74+
}
6975
if (param.transportPrimaryOnnxCCDBPath.empty()) {
7076
throw std::runtime_error("Stack.transportPrimaryOnnxCCDBPath must be configured");
7177
}
@@ -94,63 +100,36 @@ class OnnxPrimaryTransport
94100
mModel.initSessionFromBuffer(mModelBytes.data(), mModelBytes.size());
95101

96102
const auto inputShapes = mModel.getNumInputNodes();
97-
if (inputShapes.size() != 1 || inputShapes[0].empty() ||
98-
(inputShapes[0].back() > 0 && inputShapes[0].back() != OnnxFeatureCount)) {
103+
if (inputShapes.size() != 1 || inputShapes[0].size() != 2 ||
104+
(inputShapes[0][0] != 1 && inputShapes[0][0] != -1) ||
105+
inputShapes[0][1] != OnnxFeatureCount) {
99106
throw std::runtime_error("primary transport ONNX model must have one float input with 25 features");
100107
}
101-
if (mModel.getNumOutputNodes().size() != 1 || mOutputIndex < 0) {
102-
throw std::runtime_error("primary transport ONNX model must have one output and a non-negative output index");
108+
const auto outputShapes = mModel.getNumOutputNodes();
109+
if (outputShapes.size() != 1 || outputShapes[0].size() != 2 ||
110+
(outputShapes[0][0] != 1 && outputShapes[0][0] != -1) ||
111+
outputShapes[0][1] <= 0 || mOutputIndex < 0 || mOutputIndex >= outputShapes[0][1]) {
112+
throw std::runtime_error("track transport ONNX model must have one [batch, scores] output and a valid score index");
103113
}
104114
}
105115

106-
bool transport(const TParticle& particle, const std::vector<TParticle>& primaries)
116+
bool transport(const TParticle& particle, double eventX, double eventY, double eventZ)
107117
{
108-
std::vector<std::vector<float>> inputs{makeFeatures(particle, primaries)};
118+
// OrtModel mutates its shape buffers during inference. Stack clones share
119+
// this classifier, so protect the session and those buffers together.
120+
std::lock_guard<std::mutex> lock(mMutex);
121+
std::vector<std::vector<float>> inputs{o2::data::detail::makeTrackTransportFeatures(particle, eventX, eventY, eventZ)};
109122
auto output = mModel.inference<float, float>(inputs);
110123
if (static_cast<size_t>(mOutputIndex) >= output.size()) {
111124
throw std::runtime_error("Stack.transportPrimaryOnnxOutputIndex is outside the model output");
112125
}
113-
float score = output[mOutputIndex];
114-
if (mApplySigmoid) {
115-
score = score >= 0.f ? 1.f / (1.f + std::exp(-score)) : std::exp(score) / (1.f + std::exp(score));
116-
}
117-
// Class 1 means that GEANT transport can be avoided.
118-
return score < mThreshold;
126+
return o2::data::detail::transportFromOnnxScore(output[mOutputIndex], mThreshold, mApplySigmoid);
119127
}
120128

121129
private:
122-
static std::vector<float> makeFeatures(const TParticle& particle, const std::vector<TParticle>& primaries)
123-
{
124-
const double px = particle.Px();
125-
const double py = particle.Py();
126-
const double pz = particle.Pz();
127-
const double momentum = std::sqrt(px * px + py * py + pz * pz);
128-
const double pt = std::hypot(px, py);
129-
const double mass = particle.GetMass();
130-
const double energy = std::sqrt(std::max(0., mass * mass + momentum * momentum));
131-
const double eta = momentum > std::abs(pz) ? 0.5 * std::log((momentum + pz) / (momentum - pz)) : 0.;
132-
const double theta = momentum > 0. ? std::acos(pz / momentum) : 0.;
133-
const double rapidity = energy > std::abs(pz) ? 0.5 * std::log((energy + pz) / (energy - pz)) : 0.;
134-
const auto* pdgInfo = particle.GetPDG();
135-
const double chargeSign = pdgInfo == nullptr || pdgInfo->Charge() == 0. ? 0. : std::copysign(1., pdgInfo->Charge());
136-
const TParticle& eventReference = primaries.empty() ? particle : primaries.front();
137-
const double dx = particle.Vx() - eventReference.Vx();
138-
const double dy = particle.Vy() - eventReference.Vy();
139-
const double dz = particle.Vz() - eventReference.Vz();
140-
const double pdg = particle.GetPdgCode();
141-
142-
return {static_cast<float>(pdg), static_cast<float>(std::abs(pdg)), static_cast<float>(chargeSign),
143-
static_cast<float>(mass), static_cast<float>(energy), static_cast<float>(energy - mass),
144-
static_cast<float>(px), static_cast<float>(py), static_cast<float>(pz), static_cast<float>(momentum),
145-
static_cast<float>(pt), static_cast<float>(eta), static_cast<float>(particle.Phi()), static_cast<float>(theta),
146-
static_cast<float>(rapidity), static_cast<float>(particle.Vx()), static_cast<float>(particle.Vy()),
147-
static_cast<float>(particle.Vz()), static_cast<float>(particle.T()), static_cast<float>(dx),
148-
static_cast<float>(dy), static_cast<float>(dz), static_cast<float>(std::hypot(particle.Vx(), particle.Vy())),
149-
static_cast<float>(std::hypot(dx, dy)), static_cast<float>(std::sqrt(dx * dx + dy * dy + dz * dz))};
150-
}
151-
130+
std::vector<char> mModelBytes; // Must outlive the session (members are destroyed in reverse order).
152131
o2::ml::OrtModel mModel;
153-
std::vector<char> mModelBytes; // ORT may use model bytes directly; retain them for the session lifetime.
132+
std::mutex mMutex;
154133
float mThreshold;
155134
int mOutputIndex;
156135
bool mApplySigmoid;
@@ -221,20 +200,24 @@ Stack::Stack(Int_t size)
221200
LOG(info) << "Successfully retrieve external \'transportPrimary\' frunction: " << param.transportPrimaryFileName;
222201
} else if (param.transportPrimary.compare("onnx") == 0) {
223202
try {
224-
auto classifier = std::make_shared<OnnxPrimaryTransport>(param);
225-
transportPrimary = [classifier](const TParticle& p, const std::vector<TParticle>& particles) {
226-
return classifier->transport(p, particles);
203+
auto classifier = std::make_shared<OnnxTrackTransport>(param);
204+
mTransportTrack = [classifier, invert = param.transportPrimaryInvert](const TParticle& p, double x, double y, double z) {
205+
const bool transport = classifier->transport(p, x, y, z);
206+
return invert ? !transport : transport;
227207
};
228-
LOG(info) << "Successfully configured ONNX primary transport pruning from CCDB path "
208+
// ONNX runs at PreTrack, where the true event vertex and both primary
209+
// and secondary birth states are available, also in parallel simulation.
210+
transportPrimary = [](const TParticle&, const std::vector<TParticle>&) { return true; };
211+
LOG(info) << "Successfully configured ONNX track transport pruning from CCDB path "
229212
<< param.transportPrimaryOnnxCCDBPath;
230213
} catch (const std::exception& error) {
231-
LOG(fatal) << "Failed to configure ONNX primary transport pruning: " << error.what();
214+
LOG(fatal) << "Failed to configure ONNX track transport pruning: " << error.what();
232215
}
233216
} else {
234217
LOG(fatal) << "unsupported \'trasportPrimary\' mode: " << param.transportPrimary;
235218
}
236219

237-
if (param.transportPrimaryInvert) {
220+
if (param.transportPrimaryInvert && param.transportPrimary != "onnx") {
238221
mTransportPrimary = [transportPrimary](const TParticle& p, const std::vector<TParticle>& particles) { return !transportPrimary(p, particles); };
239222
} else {
240223
mTransportPrimary = transportPrimary;
@@ -259,7 +242,9 @@ Stack::Stack(const Stack& rhs)
259242
mMinHits(rhs.mMinHits),
260243
mEnergyCut(rhs.mEnergyCut),
261244
mTrackRefs(new std::vector<o2::TrackReference>),
262-
mIsG4Like(rhs.mIsG4Like)
245+
mIsG4Like(rhs.mIsG4Like),
246+
mTransportPrimary(rhs.mTransportPrimary),
247+
mTransportTrack(rhs.mTransportTrack)
263248
{
264249
LOG(debug) << "copy constructor called";
265250
mTracks = new std::vector<MCTrack>();
@@ -295,6 +280,8 @@ Stack& Stack::operator=(const Stack& rhs)
295280
mMinHits = rhs.mMinHits;
296281
mEnergyCut = rhs.mEnergyCut;
297282
mIsG4Like = rhs.mIsG4Like;
283+
mTransportPrimary = rhs.mTransportPrimary;
284+
mTransportTrack = rhs.mTransportTrack;
298285

299286
return *this;
300287
}
@@ -395,6 +382,34 @@ void Stack::handleTransportPrimary(TParticle& p)
395382
}
396383
}
397384

385+
bool Stack::transportTrack(const TParticle& particle, double eventX, double eventY, double eventZ)
386+
{
387+
if (!mTransportTrack || mTransportTrack(particle, eventX, eventY, eventZ)) {
388+
return true;
389+
}
390+
// Keep bookkeeping and ancestry, but mark the birth track as inhibited.
391+
// Do not alter the primary completion count: its PreTrack/FinishPrimary
392+
// lifecycle has already started and will still be completed by the engine.
393+
auto inhibit = [](TParticle& p) {
394+
p.SetBit(ParticleStatus::kToBeDone, 0);
395+
p.SetBit(ParticleStatus::kInhibited, 1);
396+
};
397+
inhibit(mCurrentParticle);
398+
const int id = mIndexOfCurrentTrack;
399+
if (id >= 0 && id < static_cast<int>(mPrimaryParticles.size())) {
400+
inhibit(mPrimaryParticles[id]);
401+
(*mTracks)[id].setToBeDone(false);
402+
(*mTracks)[id].setInhibited(true);
403+
} else if (id >= 0 && id < static_cast<int>(mTrackIDtoParticlesEntry.size())) {
404+
const int entry = mTrackIDtoParticlesEntry[id];
405+
if (entry >= 0 && entry < static_cast<int>(mParticles.size())) {
406+
mParticles[entry].setToBeDone(false);
407+
mParticles[entry].setInhibited(true);
408+
}
409+
}
410+
return false;
411+
}
412+
398413
void Stack::PushTrack(int toBeDone, TParticle& p)
399414
{
400415
// printf("stack -> Pushing Primary toBeDone %5d %5d parentId %5d pdgCode %5d is %5d entries %5d \n", toBeDone, p.TestBit(ParticleStatus::kToBeDone), p.GetFirstMother(), p.GetPdgCode(), p.GetStatusCode(), mNumberOfEntriesInParticles);

‎Detectors/Base/test/testStack.cxx‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#include <boost/test/unit_test.hpp>
1616
#include "DetectorsBase/Detector.h"
1717
#include "DetectorsBase/Stack.h"
18+
#include "DetectorsBase/TrackTransportUtils.h"
1819
#include "SimulationDataFormat/BaseHits.h"
1920
#include "TFile.h"
2021
#include "TMCProcess.h"
@@ -180,3 +181,37 @@ BOOST_AUTO_TEST_CASE(Offsetting_keeps_an_invalid_index_invalid)
180181
BOOST_CHECK_EQUAL(o2::base::Detector::offsetTrackIndex(7, nprimaries, primaryOffset, secondaryOffset), 107);
181182
BOOST_CHECK_EQUAL(o2::base::Detector::offsetTrackIndex(-1, nprimaries, primaryOffset, secondaryOffset), -1);
182183
}
184+
185+
BOOST_AUTO_TEST_CASE(Track_transport_features_match_training_units)
186+
{
187+
// A secondary displaced from the actual event vertex, with negative phi.
188+
TParticle p(211, 0, 0, -1, -1, -1, 0., -2., 0., 2.1, 11., 22., 33., 7.e-9);
189+
auto f = o2::data::detail::makeTrackTransportFeatures(p, 1., 2., 3.);
190+
BOOST_REQUIRE_EQUAL(f.size(), 25);
191+
BOOST_CHECK_CLOSE(f[18], 7.f, 1.e-4f); // nanoseconds, not seconds
192+
BOOST_CHECK_CLOSE(f[12], -std::acos(-1.f) / 2.f, 1.e-4f); // atan2 range
193+
BOOST_CHECK_EQUAL(f[19], 10.f);
194+
BOOST_CHECK_EQUAL(f[20], 20.f);
195+
BOOST_CHECK_EQUAL(f[21], 30.f);
196+
BOOST_CHECK_CLOSE(f[23], std::sqrt(500.f), 1.e-4f);
197+
BOOST_CHECK_EQUAL(f[2], 1.f);
198+
// Undefined angular inputs retain CSV missing-value semantics.
199+
p.SetMomentum(0., 0., 0., 0.);
200+
f = o2::data::detail::makeTrackTransportFeatures(p, 1., 2., 3.);
201+
BOOST_CHECK(std::isnan(f[11]));
202+
BOOST_CHECK(std::isnan(f[13]));
203+
}
204+
205+
BOOST_AUTO_TEST_CASE(Track_transport_class_one_rejects_and_invalid_scores_fail)
206+
{
207+
using o2::data::detail::transportFromOnnxScore;
208+
BOOST_CHECK(transportFromOnnxScore(0.1f, 0.5f, false));
209+
BOOST_CHECK(!transportFromOnnxScore(0.9f, 0.5f, false));
210+
BOOST_CHECK(!transportFromOnnxScore(0.f, 0.5f, true));
211+
BOOST_CHECK(transportFromOnnxScore(-1000.f, 0.5f, true));
212+
BOOST_CHECK(!transportFromOnnxScore(1000.f, 0.5f, true));
213+
BOOST_CHECK_THROW(transportFromOnnxScore(std::numeric_limits<float>::quiet_NaN(), 0.5f, false), std::runtime_error);
214+
BOOST_CHECK_THROW(transportFromOnnxScore(std::numeric_limits<float>::infinity(), 0.5f, true), std::runtime_error);
215+
BOOST_CHECK_THROW(transportFromOnnxScore(2.f, 0.5f, false), std::runtime_error);
216+
BOOST_CHECK_THROW(transportFromOnnxScore(0.5f, -1.f, false), std::runtime_error);
217+
}

0 commit comments

Comments
 (0)