Skip to content

[NN CF] Momentum vector usage by neural network cluster finder #15582

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ChSonnabend wants to merge 3 commits into
base: dev
Choose a base branch
from
Draft

[NN CF] Momentum vector usage by neural network cluster finder #15582

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ea4d277
Momentum vector usage created by neural network
ChSonnabend Jul 6, 2026
f97e094
Please consider the following formatting changes
alibuild Jul 6, 2026
2efd2f2
Merge pull request #43 from alibuild/alibot-cleanup-15582
ChSonnabend Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,9 @@ struct ClusterNative {
static constexpr int scaleSaturatedQtot = 8;
static constexpr int maxRegularQtot = 25 * 1024;
static constexpr int maxSaturatedQtot = (USHRT_MAX - maxRegularQtot) * scaleSaturatedQtot;
static constexpr float scaleNNDirectionPacked = 2048.f;
static constexpr int16_t invalidNNDirectionPacked = -32768;
static constexpr float maxNNDirection = 32767.f / scaleNNDirectionPacked;

uint32_t timeFlagsPacked; //< Contains the time in the lower 24 bits in a packed format, contains the flags in the
// upper 8 bits
Expand All @@ -74,18 +77,40 @@ struct ClusterNative {
uint8_t sigmaPadPacked; //< Sigma of the pad in packed format
uint16_t qMax; //< QMax of the cluster
uint16_t qTotPacked; //< Total charge of the cluster
int16_t nnDydxPacked = invalidNNDirectionPacked;
int16_t nnDzdxPacked = invalidNNDirectionPacked;

GPUd() static uint16_t packPad(float pad) { return (uint16_t)(pad * scalePadPacked + 0.5); }
GPUd() static uint32_t packTime(float time) { return (uint32_t)(time * scaleTimePacked + 0.5); }
GPUd() static float unpackPad(uint16_t pad) { return float(pad) * (1.f / scalePadPacked); }
GPUd() static float unpackTime(uint32_t time) { return float(time) * (1.f / scaleTimePacked); }
GPUd() static int16_t packNNDirection(float v)
{
v = v < -maxNNDirection ? -maxNNDirection : (v > maxNNDirection ? maxNNDirection : v);
return static_cast<int16_t>(v * scaleNNDirectionPacked + (v >= 0.f ? 0.5f : -0.5f));
}
GPUd() static float unpackNNDirection(int16_t v) { return float(v) * (1.f / scaleNNDirectionPacked); }

GPUdDefault() ClusterNative() = default;
GPUd() ClusterNative(uint32_t time, uint8_t flags, uint16_t pad, uint8_t sigmaTime, uint8_t sigmaPad, uint16_t qmax, uint16_t qtotPacked) : padPacked(pad), sigmaTimePacked(sigmaTime), sigmaPadPacked(sigmaPad), qMax(qmax), qTotPacked(qtotPacked)
{
setTimePackedFlags(time, flags);
}

GPUd() bool hasNNDirection() const { return nnDydxPacked != invalidNNDirectionPacked && nnDzdxPacked != invalidNNDirectionPacked; }
GPUd() float getNNDydx() const { return unpackNNDirection(nnDydxPacked); }
GPUd() float getNNDzdx() const { return unpackNNDirection(nnDzdxPacked); }
GPUd() void setNNDirection(float dydx, float dzdx)
{
nnDydxPacked = packNNDirection(dydx);
nnDzdxPacked = packNNDirection(dzdx);
}
GPUd() void clearNNDirection()
{
nnDydxPacked = invalidNNDirectionPacked;
nnDzdxPacked = invalidNNDirectionPacked;
}

GPUd() uint16_t getQmax() const { return qMax; }
GPUd() uint32_t getQtot() const { return isSaturated() ? getSaturatedQtot() : (uint32_t)qTotPacked; }
GPUd() uint8_t getFlags() const { return timeFlagsPacked >> 24; }
Expand Down Expand Up @@ -187,6 +212,10 @@ struct ClusterNative {
return (this->qMax < rhs.qMax);
} else if (this->qTotPacked != rhs.qTotPacked) {
return (this->getQtot() < rhs.getQtot());
} else if (this->nnDydxPacked != rhs.nnDydxPacked) {
return (this->nnDydxPacked < rhs.nnDydxPacked);
} else if (this->nnDzdxPacked != rhs.nnDzdxPacked) {
return (this->nnDzdxPacked < rhs.nnDzdxPacked);
} else {
return (this->getFlags() < rhs.getFlags());
}
Expand All @@ -200,6 +229,8 @@ struct ClusterNative {
this->sigmaPadPacked == rhs.sigmaPadPacked &&
this->qMax == rhs.qMax &&
this->qTotPacked == rhs.qTotPacked &&
this->nnDydxPacked == rhs.nnDydxPacked &&
this->nnDzdxPacked == rhs.nnDzdxPacked &&
this->getFlags() == rhs.getFlags();
}

Expand Down
5 changes: 5 additions & 0 deletions GPU/GPUTracking/Base/GPUConstantMem.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,11 @@ GPUdi() GPUconstantref() const GPUParam& GPUProcessor::Param() const
return GetConstantMem()->param;
}

GPUdi() const o2::tpc::ClusterNativeAccess* GPUProcessor::GetClustersNative() const
{
return GetConstantMem()->ioPtrs.clustersNative;
}

GPUdi() void GPUProcessor::raiseError(uint32_t code, uint32_t param1, uint32_t param2, uint32_t param3) const
{
GetConstantMem()->errorCodes.raiseError(code, param1, param2, param3);
Expand Down
6 changes: 6 additions & 0 deletions GPU/GPUTracking/Base/GPUProcessor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,11 @@
#include <algorithm>
#endif

namespace o2::tpc

@davidrohr davidrohr Jul 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why does this have to go here? This is one of the base classes, included almost everywhere, so I would try to keep it as slim as possible.

{
struct ClusterNativeAccess;
}

namespace o2::gpu
{
struct GPUTrackingInOutPointers;
Expand All @@ -49,6 +54,7 @@ class GPUProcessor
#endif

GPUd() GPUconstantref() const GPUConstantMem* GetConstantMem() const; // Body in GPUConstantMem.h to avoid circular headers
GPUd() const o2::tpc::ClusterNativeAccess* GetClustersNative() const; // ...
GPUd() GPUconstantref() const GPUParam& Param() const; // ...
GPUd() void raiseError(uint32_t code, uint32_t param1 = 0, uint32_t param2 = 0, uint32_t param3 = 0) const;
const GPUReconstruction& GetRec() const { return *mRec; }
Expand Down
2 changes: 1 addition & 1 deletion GPU/GPUTracking/Definitions/GPUSettingsList.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,7 +295,7 @@ AddOption(nnLoadFromCCDB, int, 0, "", 0, "If 1 networks are fetched from ccdb, e
AddOption(nnCCDBDumpToFile, int, 0, "", 0, "If 1, additionally dump fetched CCDB networks to nnLocalFolder")
AddOption(nnLocalFolder, std::string, ".", "", 0, "Local folder in which the networks will be fetched")
AddOption(nnCCDBPath, std::string, "Users/c/csonnabe/TPC/Clusterization", "", 0, "Folder path containing the networks")
AddOption(nnCCDBWithMomentum, std::string, "", "", 0, "Distinguishes between the network with and without momentum output for the regression")
AddOption(nnCCDBWithMomentum, std::string, "0", "", 0, "Distinguishes between the network with and without momentum output for the regression")
AddOption(nnCCDBClassificationLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")
AddOption(nnCCDBRegressionLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")
AddOption(nnCCDBBeamType, std::string, "pp", "", 0, "Distinguishes between networks trained for different beam types. Options: pp, pPb, PbPb")
Expand Down
10 changes: 6 additions & 4 deletions GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,15 +120,17 @@ GPUdii() void GPUTPCNeighboursFinder::Thread<0>(int32_t /*nBlocks*/, int32_t nTh
const GPUglobalref() cahit2& hitData = pHitData[lHitNumberOffset + ih];
const float y = y0 + hitData.x * stepY;
const float z = z0 + hitData.y * stepZ;
float nnDydx = 0.f, nnDzdx = 0.f;
const bool useNNDir = tracker.HitNNDirection(row, ih, nnDydx, nnDzdx) && CAMath::Abs(nnDydx) < 10.f && CAMath::Abs(nnDzdx) < 10.f;

@davidrohr davidrohr Jul 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

there should be an option in param.rec.tpc, that enables / disables it, and this should be ANDed as first check here.


uint32_t nNeighUp = 0;
float minZ, maxZ, minY, maxY;
int32_t binYmin, binYmax, binZmin, binZmax;
int32_t nY;

{ // area in the upper row
const float yy = y * s.mUpTx;
const float zz = z * kAreaSlopeZUp;
const float yy = useNNDir ? y + nnDydx * s.mUpDx : y * s.mUpTx;
const float zz = useNNDir ? z + nnDzdx * s.mUpDx : z * kAreaSlopeZUp;
minZ = zz - kAreaSizeZUp;
maxZ = zz + kAreaSizeZUp;
minY = yy - kAreaSizeY;
Expand Down Expand Up @@ -196,8 +198,8 @@ GPUdii() void GPUTPCNeighboursFinder::Thread<0>(int32_t /*nBlocks*/, int32_t nTh
}

{ // area in the lower row
const float yy = y * s.mDnTx;
const float zz = z * kAreaSlopeZDn;
const float yy = useNNDir ? y + nnDydx * s.mDnDx : y * s.mDnTx;
const float zz = useNNDir ? z + nnDzdx * s.mDnDx : z * kAreaSlopeZDn;
minZ = zz - kAreaSizeZDn;
maxZ = zz + kAreaSizeZDn;
minY = yy - kAreaSizeY;
Expand Down
17 changes: 17 additions & 0 deletions GPU/GPUTracking/SectorTracker/GPUTPCTracker.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
#include "GPUTPCTrackParam.h"
#include "GPUTPCTracklet.h"
#include "GPUProcessor.h"
#include "DataFormatsTPC/ClusterNative.h"

namespace o2::gpu
{
Expand Down Expand Up @@ -151,6 +152,22 @@ class GPUTPCTracker : public GPUProcessor

GPUhd() int32_t HitInputID(const GPUTPCRow& row, int32_t hitIndex) const { return mData.ClusterDataIndex(row, hitIndex); }

GPUd() bool HitNNDirection(const GPUTPCRow& row, int32_t hitIndex, float& dydx, float& dzdx) const
{
const o2::tpc::ClusterNativeAccess* native = GetClustersNative();
if (native == nullptr || native->clustersLinear == nullptr) {
return false;
}
const uint32_t sectorOffset = native->clusterOffset[mISector][0];
const o2::tpc::ClusterNative& cluster = native->clustersLinear[sectorOffset + mData.ClusterDataIndex(row, hitIndex)];
if (!cluster.hasNNDirection()) {
return false;
}
dydx = cluster.getNNDydx();
dzdx = cluster.getNNDzdx();
return true;
}

/**
* The hit weight is used to determine whether a hit belongs to a certain tracklet or another one
* competing for the same hit. The tracklet that has a higher weight wins. Comparison is done
Expand Down
11 changes: 9 additions & 2 deletions GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,9 +170,16 @@ GPUdic(2, 1) void GPUTPCTrackletConstructor::UpdateTracklet(int32_t /*nBlocks*/,

float ri = 1.f / CAMath::Sqrt(dx * dx + dy * dy);
if (iRow == (int32_t)r.mStartRow + 2) {
tParam.SetSinPhi(dy * ri);
float nnDydx = 0.f, nnDzdx = 0.f;
if (tracker.HitNNDirection(row, seedIH, nnDydx, nnDzdx) && CAMath::Abs(nnDydx) < 10.f && CAMath::Abs(nnDzdx) < 10.f) {

@davidrohr davidrohr Jul 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't fully understand the math without checking in detail. You are weighting the NN direction and the slope between the clusters? That makes sense. But in any case, as first check, it should be ANDed with a setting in param.rec.tpc

@ChSonnabend ChSonnabend Jul 8, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I normalize the NN direction by the transverse element (ds = sqrt(dx2 + dy2), which becomes dx*invsqrt(1 + (dy/dx)**2)). Since sin(phi) = dy/ds you get nnDydx * nnNormI. The track parameters use the transverse path length with sin(phi) = Py / Pt and DzDs = Pz / Pt. Therefore the normalization is reused for the z direction too.

const float nnNormI = CAMath::InvSqrt(1.f + nnDydx * nnDydx);
tParam.SetSinPhi(0.5f * (dy * ri + nnDydx * nnNormI));
tParam.SetDzDs(0.5f * (dz * ri + nnDzdx * nnNormI));
} else {
tParam.SetSinPhi(dy * ri);
tParam.SetDzDs(dz * ri);
}
tParam.SetSignCosPhi(dx);
tParam.SetDzDs(dz * ri);
float err2Y, err2Z;
tracker.GetErrors2Seeding(iRow, tParam, -1.f, err2Y, err2Z); // Use correct time
tParam.SetCov(0, err2Y);
Expand Down
1 change: 1 addition & 0 deletions GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ class GPUTPCNNClusterizer : public GPUProcessor
int32_t mNnClusterizerBoundaryFillValue = -1;
int32_t mNnClusterizerModelClassNumOutputNodes = -1;
int32_t mNnClusterizerModelReg1NumOutputNodes = -1;
int32_t mNnClusterizerUseMomentumVector = 0;
int32_t mNnClusterizerModelReg2NumOutputNodes = -1;
int32_t mNnInferenceInputDType = 0; // 0: float16, 1: float32
int32_t mNnInferenceOutputDType = 0; // 0: float16, 1: float32
Expand Down
3 changes: 3 additions & 0 deletions GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,9 @@ void GPUTPCNNClusterizerHost::initClusterizer(const GPUSettingsProcessingNNclust
if (!settings.nnClusterizerUseCfRegression) {
if (mModelClass.getNumOutputNodes()[0][1] == 1 || !mModelReg2.isInitialized()) {
clustererNN.mNnClusterizerModelReg1NumOutputNodes = mModelReg1.getNumOutputNodes()[0][1];
if (clustererNN.mNnClusterizerModelReg1NumOutputNodes > 5) {
clustererNN.mNnClusterizerUseMomentumVector = true;
}
} else {
clustererNN.mNnClusterizerModelReg1NumOutputNodes = mModelReg1.getNumOutputNodes()[0][1];
clustererNN.mNnClusterizerModelReg2NumOutputNodes = mModelReg2.getNumOutputNodes()[0][1];
Expand Down
14 changes: 14 additions & 0 deletions GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,15 @@ using namespace o2::gpu::tpccf;

static_assert(GPUTPCNNClusterizerKernels::SCRATCH_PAD_WORK_GROUP_SIZE == GPUTPCCFClusterizer::SCRATCH_PAD_WORK_GROUP_SIZE, "Work group sizes do not match");

GPUd() static void setClass1NNDirection(const GPUTPCNNClusterizer& nn, int8_t dtype, int32_t base, o2::tpc::ClusterNative& cluster)
{
if (dtype == 0) {
cluster.setNNDirection(nn.mOutputDataReg1_32[base + 5], nn.mOutputDataReg1_32[base + 5]);
} else {
cluster.setNNDirection(nn.mOutputDataReg1_16[base + 6].ToFloat(), nn.mOutputDataReg1_16[base + 6].ToFloat());
}
}

// Defining individual thread functions for data filling, determining the class label and running the CF clusterizer
template <>
GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::runCfClusterizer>(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUSharedMemory& smem, processorType& processors, uint8_t sector, int8_t dtype, int8_t withMC, uint32_t batchStart)
Expand Down Expand Up @@ -476,6 +485,9 @@ GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::pub
}
return;
}
if (clustererNN.mNnClusterizerUseMomentumVector) {
setClass1NNDirection(clustererNN, dtype, model_output_index, myCluster);
}

uint32_t rowIndex = 0;
if (clusterOut != nullptr) {
Expand Down Expand Up @@ -593,6 +605,7 @@ GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::pub
}
return;
}
// setClass2NNDirection(clustererNN, dtype, model_output_index, myCluster);

uint32_t rowIndex = 0;
if (clusterOut != nullptr) {
Expand Down Expand Up @@ -646,6 +659,7 @@ GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::pub
}
return;
}
// setClass2NNDirection(clustererNN, dtype, model_output_index, myCluster);

if (clusterOut != nullptr) {
rowIndex = GPUTPCCFClusterizer::sortIntoBuckets(
Expand Down
3 changes: 2 additions & 1 deletion GPU/Workflow/src/GPUWorkflowSpec.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1259,7 +1259,6 @@ Inputs GPURecoWorkflowSpec::inputs()
std::map<std::string, std::string> metadata;
metadata["inputDType"] = nnClusterizerSettings.nnInferenceInputDType; // FP16 or FP32
metadata["outputDType"] = nnClusterizerSettings.nnInferenceOutputDType; // FP16 or FP32
metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum; // 0, 1 -> Only for regression model
metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBClassificationLayerType; // FC, CNN
metadata["nnCCDBInteractionRate"] = nnClusterizerSettings.nnCCDBInteractionRate; // in kHz
metadata["nnCCDBBeamType"] = nnClusterizerSettings.nnCCDBBeamType; // pp, pPb, PbPb
Expand All @@ -1286,6 +1285,7 @@ Inputs GPURecoWorkflowSpec::inputs()
printSettings(metadata);
}

metadata["nnCCDBWithMomentum"] = "0";
if (mSpecConfig.nnEvalMode[0] == "c1") {
metadata["nnCCDBEvalType"] = "classification_c1";
convert_map_to_metadata(metadata, ccdb_metadata);
Expand All @@ -1297,6 +1297,7 @@ Inputs GPURecoWorkflowSpec::inputs()
inputs.emplace_back("nn_classification_c2", gDataOriginTPC, "NNCLUSTERIZER_C2", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath + "/" + metadata["nnCCDBEvalType"], ccdb_metadata, 0));
}

metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum; // 0, 1 -> Only for regression model
metadata["nnCCDBEvalType"] = "regression_c1";
metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBRegressionLayerType;
convert_map_to_metadata(metadata, ccdb_metadata);
Expand Down