T3 DNN Draft PR Reference

RecoTracker/LSTCore/interface/alpaka/Common.h

@@ -69,18 +69,27 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
       {0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.18f, 0.18f, /*10*/ 0.18f, 0.18f, 0.18f, 0.18f, 0.18f},
       {0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.4f, 0.18f, /*10*/ 0.18f, 0.18f, 0.18f, 0.18f, 0.18f}};
 
+  // Common constants used by both DNNs
+  ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kEta_norm = 2.5f;
+  ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kPhi_norm = kPi;
+  constexpr unsigned int kPtBins = 2;
+  constexpr unsigned int kEtaBins = 10;
+
+  namespace t3dnn {
+    ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kZ_max = 224.149505f;
+    ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kR_max = 98.932365f;
+    ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kWp[kPtBins][kEtaBins] = {
+        {0.0387, 0.0395, 0.0599, 0.083, 0.1433, 0.163, 0.2927, 0.3243, 0.638, 0.7512},
+        {0.0019, 0.0024, 0.0047, 0.0107, 0.0107, 0.0316, 0.2533, 0.1114, 0.0501, 0.0407}};
+  } // namespace t3dnn
+
   namespace t5dnn {
     ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kZ_max = 267.2349854f;
     ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kR_max = 110.1099396f;
-    ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kEta_norm = 2.5f;
-    ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kPhi_norm = kPi;
-    // pt, eta binned
-    constexpr unsigned int kPtBins = 2;
-    constexpr unsigned int kEtaBins = 10;
     ALPAKA_STATIC_ACC_MEM_GLOBAL constexpr float kWp[kPtBins][kEtaBins] = {
         {0.4493, 0.4939, 0.5715, 0.6488, 0.5709, 0.5938, 0.7164, 0.7565, 0.8103, 0.8593},
         {0.4488, 0.4448, 0.5067, 0.5929, 0.4836, 0.4112, 0.4968, 0.4403, 0.5597, 0.5067}};
-  }  // namespace t5dnn
+  } // namespace t5dnn
 
 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE::lst
 #endif

RecoTracker/LSTCore/src/alpaka/NeuralNetworkWeights.h → RecoTracker/LSTCore/src/alpaka/T5NeuralNetworkWeights.h

@@ -1,5 +1,5 @@
-#ifndef RecoTracker_LSTCore_src_alpaka_NeuralNetworkWeights_h
-#define RecoTracker_LSTCore_src_alpaka_NeuralNetworkWeights_h
+#ifndef RecoTracker_LSTCore_src_alpaka_T5NeuralNetworkWeights_h
+#define RecoTracker_LSTCore_src_alpaka_T5NeuralNetworkWeights_h
 
 #include <alpaka/alpaka.hpp>
 

RecoTracker/LSTCore/src/alpaka/Triplet.h

@@ -13,6 +13,8 @@
 #include "MiniDoublet.h"
 #include "Hit.h"
 
+#include "NeuralNetwork.h"
+
 namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
 
   ALPAKA_FN_ACC ALPAKA_FN_INLINE void addTripletToMemory(ModulesConst modules,
@@ -730,6 +732,16 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
                                    ptCut))
       return false;
 
+    bool inference = lst::t3dnn::runInference(acc,
+                                             mds,
+                                             firstMDIndex,
+                                             secondMDIndex,
+                                             thirdMDIndex,
+                                             circleRadius,
+                                             betaIn);
+    if (!inference)  // T3-building cut
+      return false;
+
     return true;
   }
 

RecoTracker/LSTCore/standalone/code/core/AccessHelper.cc

@@ -207,10 +207,26 @@ std::vector<unsigned int> getModuleIdxsFromT5(LSTEvent* event, unsigned int T5)
   }
   return module_idxs;
 }
+
+//____________________________________________________________________________________________
+std::vector<unsigned int> getModuleIdxsFromT3(LSTEvent* event, unsigned int T3) {
+  std::vector<unsigned int> hits = getHitsFromT3(event, T3);
+  std::vector<unsigned int> module_idxs;
+  auto hitsEvt = event->getHits<HitsSoA>();
+  for (auto& hitIdx : hits) {
+    module_idxs.push_back(hitsEvt.moduleIndices()[hitIdx]);
+  }
+  return module_idxs;
+}
+
 //____________________________________________________________________________________________
 std::vector<unsigned int> getHitTypesFromT5(LSTEvent* event, unsigned int T5) {
   return {4, 4, 4, 4, 4, 4, 4, 4, 4, 4};
-  ;
+}
+
+//____________________________________________________________________________________________
+std::vector<unsigned int> getHitTypesFromT3(LSTEvent* event, unsigned int T5) {
+  return {4, 4, 4, 4, 4, 4};
 }
 
 //____________________________________________________________________________________________

RecoTracker/LSTCore/standalone/code/core/AccessHelper.h

@@ -33,6 +33,8 @@ std::tuple<std::vector<unsigned int>, std::vector<unsigned int>> getHitIdxsAndHi
 std::vector<unsigned int> getLSsFromT3(LSTEvent* event, unsigned int T3);
 std::vector<unsigned int> getMDsFromT3(LSTEvent* event, unsigned int T3);
 std::vector<unsigned int> getHitsFromT3(LSTEvent* event, unsigned int T3);
+std::vector<unsigned int> getHitTypesFromT3(LSTEvent* event, unsigned int T3);
+std::vector<unsigned int> getModuleIdxsFromT3(LSTEvent* event, unsigned int T3);
 std::tuple<std::vector<unsigned int>, std::vector<unsigned int>> getHitIdxsAndHitTypesFromT3(LSTEvent* event,
                                                                                              unsigned T3);
 

RecoTracker/LSTCore/standalone/code/core/write_lst_ntuple.cc

@@ -161,6 +161,7 @@ void createOptionalOutputBranches() {
 
   // T5 DNN branches
   createT5DNNBranches();
+  createT3DNNBranches();
 
 #endif
 }
@@ -192,6 +193,41 @@ void createT5DNNBranches() {
   }
 }
 
+
+//________________________________________________________________________________________________________________________________
+void createT3DNNBranches() {
+  // Common branches for T3 properties based on TripletsSoA fields
+  ana.tx->createBranch<std::vector<float>>("t3_betaIn");
+  ana.tx->createBranch<std::vector<float>>("t3_centerX");
+  ana.tx->createBranch<std::vector<float>>("t3_centerY");
+  ana.tx->createBranch<std::vector<float>>("t3_radius");
+  ana.tx->createBranch<std::vector<bool>>("t3_partOfPT5");
+  ana.tx->createBranch<std::vector<bool>>("t3_partOfT5");
+  ana.tx->createBranch<std::vector<bool>>("t3_partOfPT3");
+  ana.tx->createBranch<std::vector<float>>("t3_pMatched");
+  ana.tx->createBranch<std::vector<float>>("t3_sim_vxy");
+  ana.tx->createBranch<std::vector<float>>("t3_sim_vz");
+
+  // Hit-specific branches (T3 has 4 hits from two segments)
+  std::vector<std::string> hitIndices = {"0", "1", "2", "3", "4", "5"};
+  std::vector<std::string> hitProperties = {"r", "x", "y", "z", "eta", "phi", "detId", "layer", "moduleType"};
+
+  for (const auto& idx : hitIndices) {
+    for (const auto& prop : hitProperties) {
+      std::string branchName = "t3_hit_" + idx + "_" + prop;
+      if (prop == "detId" || prop == "layer" || prop == "moduleType") {
+        ana.tx->createBranch<std::vector<int>>(branchName);
+      } else {
+        ana.tx->createBranch<std::vector<float>>(branchName);
+      }
+    }
+  }
+
+  // Additional metadata branches
+  ana.tx->createBranch<std::vector<int>>("t3_layer_binary");
+  ana.tx->createBranch<std::vector<std::vector<int>>>("t3_matched_simIdx");
+}
+
 //________________________________________________________________________________________________________________________________
 void createGnnNtupleBranches() {
   // Mini Doublets
@@ -340,6 +376,7 @@ void setOptionalOutputBranches(LSTEvent* event) {
   setQuintupletOutputBranches(event);
   setPixelTripletOutputBranches(event);
   setOccupancyBranches(event);
+  setT3DNNBranches(event);
   setT5DNNBranches(event);
 
 #endif
@@ -638,6 +675,42 @@ void setPixelTripletOutputBranches(LSTEvent* event) {
   ana.tx->setBranch<std::vector<int>>("pT3_isDuplicate", pT3_isDuplicate);
 }
 
+//________________________________________________________________________________________________________________________________
+void fillT3DNNBranches(LSTEvent* event, unsigned int iT3) {
+  auto hits = event->getHits<HitsSoA>();
+  auto modules = event->getModules<ModulesSoA>();
+
+  std::vector<unsigned int> hitIdx = getHitsFromT3(event, iT3);
+  std::vector<lst_math::Hit> hitObjects;
+
+  for (int i = 0; i < hitIdx.size(); ++i) {
+    unsigned int hit = hitIdx[i];
+    float x = hits.xs()[hit];
+    float y = hits.ys()[hit];
+    float z = hits.zs()[hit];
+    lst_math::Hit hitObj(x, y, z);
+    hitObjects.push_back(hitObj);
+
+    std::string idx = std::to_string(i);
+    ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_r", sqrt(x * x + y * y));
+    ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_x", x);
+    ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_y", y);
+    ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_z", z);
+    ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_eta", hitObj.eta());
+    ana.tx->pushbackToBranch<float>("t3_hit_" + idx + "_phi", hitObj.phi());
+
+    int subdet = trk.ph2_subdet()[hits.idxs()[hit]];
+    int is_endcap = subdet == 4;
+    int layer = trk.ph2_layer()[hits.idxs()[hit]] + 6 * is_endcap;
+    int detId = trk.ph2_detId()[hits.idxs()[hit]];
+    unsigned int module = hits.moduleIndices()[hit];
+
+    ana.tx->pushbackToBranch<int>("t3_hit_" + idx + "_detId", detId);
+    ana.tx->pushbackToBranch<int>("t3_hit_" + idx + "_layer", layer);
+    ana.tx->pushbackToBranch<int>("t3_hit_" + idx + "_moduleType", modules.moduleType()[module]);
+  }
+}
+
 //________________________________________________________________________________________________________________________________
 void fillT5DNNBranches(LSTEvent* event, unsigned int iT3) {
   auto hits = event->getHits<HitsSoA>();
@@ -690,6 +763,69 @@ void fillT5DNNBranches(LSTEvent* event, unsigned int iT3) {
   ana.tx->pushbackToBranch<float>("t5_t3_phi", hitObjects[0].phi());
 }
 
+void setT3DNNBranches(LSTEvent* event) {
+  auto const triplets = event->getTriplets<TripletsSoA>();
+  auto const tripletsOccupancy = event->getTriplets<TripletsOccupancySoA>();
+  auto modules = event->getModules<ModulesSoA>();
+  auto ranges = event->getRanges();
+
+  for (unsigned int lowerModuleIdx = 0; lowerModuleIdx < modules.nLowerModules(); ++lowerModuleIdx) {
+    int nTriplets = tripletsOccupancy.nTriplets()[lowerModuleIdx];
+    for (unsigned int idx = 0; idx < nTriplets; idx++) {
+      unsigned int tripletIndex = ranges.tripletModuleIndices()[lowerModuleIdx] + idx;
+
+      // Get hit indices and types
+      std::vector<unsigned int> hit_idx = getHitsFromT3(event, tripletIndex);
+      std::vector<unsigned int> hit_type = getHitTypesFromT3(event, tripletIndex);
+      std::vector<unsigned int> module_idx = getModuleIdxsFromT3(event, tripletIndex);
+
+      // Calculate layer binary representation
+      int layer_binary = 0;
+      for (size_t i = 0; i < module_idx.size(); i += 2) {
+        layer_binary |= (1 << (modules.layers()[module_idx[i]] + 6 * (modules.subdets()[module_idx[i]] == 4)));
+      }
+
+      // Get matching information with percent matched
+      float percent_matched;
+      std::vector<int> simidx = matchedSimTrkIdxs(hit_idx, hit_type, false, &percent_matched);
+
+      // Fill the branches with T3-specific data
+      ana.tx->pushbackToBranch<float>("t3_betaIn", triplets.betaIn()[tripletIndex]);
+      ana.tx->pushbackToBranch<float>("t3_centerX", triplets.centerX()[tripletIndex]);
+      ana.tx->pushbackToBranch<float>("t3_centerY", triplets.centerY()[tripletIndex]);
+      ana.tx->pushbackToBranch<float>("t3_radius", triplets.radius()[tripletIndex]);
+      ana.tx->pushbackToBranch<bool>("t3_partOfPT5", triplets.partOfPT5()[tripletIndex]);
+      ana.tx->pushbackToBranch<bool>("t3_partOfT5", triplets.partOfT5()[tripletIndex]);
+      ana.tx->pushbackToBranch<bool>("t3_partOfPT3", triplets.partOfPT3()[tripletIndex]);
+      ana.tx->pushbackToBranch<int>("t3_layer_binary", layer_binary);
+      ana.tx->pushbackToBranch<std::vector<int>>("t3_matched_simIdx", simidx);
+      ana.tx->pushbackToBranch<float>("t3_pMatched", percent_matched);
+
+      // Add vertex information for matched sim tracks
+      if (simidx.size() == 0) {
+        // No matched sim track - set default values
+        ana.tx->pushbackToBranch<float>("t3_sim_vxy", 0.0);
+        ana.tx->pushbackToBranch<float>("t3_sim_vz", 0.0);
+      } else {
+        // Get vertex information from the first matched sim track
+        int vtxidx = trk.sim_parentVtxIdx()[simidx[0]];
+        float vtx_x = trk.simvtx_x()[vtxidx];
+        float vtx_y = trk.simvtx_y()[vtxidx];
+        float vtx_z = trk.simvtx_z()[vtxidx];
+
+        // Calculate transverse distance from origin
+        float vxy = sqrt(vtx_x * vtx_x + vtx_y * vtx_y);
+
+        ana.tx->pushbackToBranch<float>("t3_sim_vxy", vxy);
+        ana.tx->pushbackToBranch<float>("t3_sim_vz", vtx_z);
+      }
+
+      // Fill hit-specific information
+      fillT3DNNBranches(event, tripletIndex);
+    }
+  }
+}
+
 //________________________________________________________________________________________________________________________________
 void setT5DNNBranches(LSTEvent* event) {
   auto triplets = event->getTriplets<TripletsOccupancySoA>();

RecoTracker/LSTCore/standalone/code/core/write_lst_ntuple.h

@@ -19,6 +19,7 @@ void createRequiredOutputBranches();
 void createOptionalOutputBranches();
 void createGnnNtupleBranches();
 void createT5DNNBranches();
+void createT3DNNBranches();
 
 void fillOutputBranches(LSTEvent* event);
 void setOutputBranches(LSTEvent* event);
@@ -30,7 +31,9 @@ void setPixelTripletOutputBranches(LSTEvent* event);
 void setGnnNtupleBranches(LSTEvent* event);
 void setGnnNtupleMiniDoublet(LSTEvent* event, unsigned int MD);
 void fillT5DNNBranches(LSTEvent* event, unsigned int T3);
+void fillT3DNNBranches(LSTEvent* event, unsigned int iT3);
 void setT5DNNBranches(LSTEvent* event);
+void setT3DNNBranches(LSTEvent* event);
 
 std::tuple<int, float, float, float, int, std::vector<int>> parseTrackCandidate(LSTEvent* event, unsigned int);
 std::tuple<float, float, float, std::vector<unsigned int>, std::vector<unsigned int>> parsepT5(LSTEvent* event,