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

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

@@ -291,10 +291,42 @@ struct ClusterNative {
   }
 };
 
+struct ClusterNativeNNDirection {
+  static constexpr float scalePacked = 2048.f;
+  static constexpr float maxDirection = 32767.f / scalePacked;
+  static constexpr uint16_t invalidPacked = 0;
+
+  uint16_t dydxPacked = invalidPacked;
+  uint16_t dzdxPacked = invalidPacked;
+
+  GPUd() static uint16_t pack(float v)
+  {
+    v = v < -maxDirection ? -maxDirection : (v > maxDirection ? maxDirection : v);
+    const int32_t packedSigned = static_cast<int32_t>(v * scalePacked + (v >= 0.f ? 0.5f : -0.5f));
+    return static_cast<uint16_t>(packedSigned + 32768);
+  }
+  GPUd() static float unpack(uint16_t v) { return static_cast<float>(static_cast<int32_t>(v) - 32768) * (1.f / scalePacked); }
+
+  GPUd() bool hasDirection() const { return dydxPacked != invalidPacked && dzdxPacked != invalidPacked; }
+  GPUd() float getDydx() const { return unpack(dydxPacked); }
+  GPUd() float getDzdx() const { return unpack(dzdxPacked); }
+  GPUd() void set(float dydx, float dzdx)
+  {
+    dydxPacked = pack(dydx);
+    dzdxPacked = pack(dzdx);
+  }
+  GPUd() void clear()
+  {
+    dydxPacked = invalidPacked;
+    dzdxPacked = invalidPacked;
+  }
+};
+
 // This is an index struct to access TPC clusters inside sectors and rows. It shall not own the data, but just point to
 // the data inside a buffer.
 struct ClusterNativeAccess {
   const ClusterNative* clustersLinear;
+  const ClusterNativeNNDirection* clustersLinearNNDirection = nullptr;
   const ClusterNative* clusters[constants::MAXSECTOR][constants::MAXGLOBALPADROW];
   const o2::dataformats::ConstMCTruthContainerView<o2::MCCompLabel>* clustersMCTruth;
   unsigned int nClusters[constants::MAXSECTOR][constants::MAXGLOBALPADROW];

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

@@ -167,6 +167,7 @@ AddOptionRTC(dEdxClusterRejectionFlagMask, int8_t, o2::gpu::GPUTPCGMMergedTrackH
 AddOptionRTC(dEdxClusterRejectionFlagMaskAlt, int8_t, o2::gpu::GPUTPCGMMergedTrackHit::flagEdge, "", 0, "OR mask of TPC flags that will reject the cluster in alternative dEdx")
 AddOptionRTC(rejectEdgeClustersInSeeding, int8_t, 0, "", 0, "Reject edge clusters based on uncorrected track Y during seeding")
 AddOptionRTC(rejectEdgeClustersInTrackFit, int8_t, 0, "", 0, "Reject edge clusters based on uncorrected track Y during track fit")
+AddOptionRTC(useNNClusterDirection, int8_t, 0, "", 0, "Use TPC NN cluster direction estimate in track seeding")
 AddOptionRTC(tubeExtraProtectMinRow, uint8_t, 20, "", 0, "Increase Protection, decrease removal by factor 2, when below this row")
 AddOptionRTC(tubeExtraProtectEdgePads, uint8_t, 2, "", 0, "Increase Protection, decrease removal by factor 2, when on this number of pads from the edge")
 
@@ -295,7 +296,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/Global/GPUChainTracking.h

@@ -41,6 +41,7 @@ namespace o2::tpc
 {
 struct ClusterNativeAccess;
 struct ClusterNative;
+struct ClusterNativeNNDirection;
 class CalibdEdxContainer;
 } // namespace o2::tpc
 
@@ -313,7 +314,7 @@ class GPUChainTracking : public GPUChain
   void RunTPCClusterFilter(o2::tpc::ClusterNativeAccess* clusters, std::function<o2::tpc::ClusterNative*(size_t)> allocator, bool applyClusterCuts);
   bool NeedTPCClustersOnGPU();
   void WriteReducedClusters();
-  void SortClusters(bool buildNativeGPU, bool propagateMCLabels, o2::tpc::ClusterNativeAccess* clusterAccess, o2::tpc::ClusterNative* clusters);
+  void SortClusters(bool buildNativeGPU, bool propagateMCLabels, o2::tpc::ClusterNativeAccess* clusterAccess, o2::tpc::ClusterNative* clusters, o2::tpc::ClusterNativeNNDirection* directions);
   template <int32_t I>
   int32_t RunTRDTrackingInternal();
   uint32_t StreamForSector(uint32_t sector) const;