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

DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h

@@ -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
@@ -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; }
@@ -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());
     }
@@ -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();
   }
 

GPU/GPUTracking/Base/GPUConstantMem.h

@@ -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);

GPU/GPUTracking/Base/GPUProcessor.h

@@ -23,6 +23,11 @@
 #include <algorithm>
 #endif
 
+namespace o2::tpc
+{
+struct ClusterNativeAccess;
+}
+
 namespace o2::gpu
 {
 struct GPUTrackingInOutPointers;
@@ -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; }

GPU/GPUTracking/Definitions/GPUSettingsList.h

@@ -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")

GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx

@@ -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;
 
     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;
@@ -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;

GPU/GPUTracking/SectorTracker/GPUTPCTracker.h

@@ -26,6 +26,7 @@
 #include "GPUTPCTrackParam.h"
 #include "GPUTPCTracklet.h"
 #include "GPUProcessor.h"
+#include "DataFormatsTPC/ClusterNative.h"
 
 namespace o2::gpu
 {
@@ -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

GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx

@@ -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) {
+            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);

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h

@@ -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

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx

@@ -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];

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx

@@ -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)
@@ -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) {
@@ -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) {
@@ -646,6 +659,7 @@ GPUdii() void GPUTPCNNClusterizerKernels::Thread<GPUTPCNNClusterizerKernels::pub
     }
     return;
   }
+  // setClass2NNDirection(clustererNN, dtype, model_output_index, myCluster);
 
   if (clusterOut != nullptr) {
     rowIndex = GPUTPCCFClusterizer::sortIntoBuckets(

GPU/Workflow/src/GPUWorkflowSpec.cxx

@@ -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
@@ -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);
@@ -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);