Project CMSSW displayed by LXR CMSSW_15_1_X_2025-04-25-2300/​CommonTools/​RecoAlgos/​interface/​KDTreeLinkerAlgo.h	Source navigation Diff markup Identifier search General search
	Version: CMSSW_15_1_X_2025-04-25-2300 CMSSW_15_1_X_2025-04-24-2300 CMSSW_15_1_X_2025-04-23-2300 CMSSW_15_1_X_2025-04-22-2300 CMSSW_15_1_X_2025-04-21-2300 CMSSW_15_1_X_2025-04-20-2300 CMSSW_15_0_4 CMSSW_15_0_3_patch1 CMSSW_14_0_21_patch1 Revert   Change

File indexing completed on 2025-04-13 22:49:35

 #ifndef KDTreeLinkerAlgoTemplated_h
 #define KDTreeLinkerAlgoTemplated_h
 
 #include "DataFormats/Math/interface/logic.h"
 
 #include <cassert>
 #include <vector>
 #include <array>
 #include <algorithm>
 
 // Box structure used to define 2D field.
 // It's used in KDTree building step to divide the detector
 // space (ECAL, HCAL...) and in searching step to create a bounding
 // box around the demanded point (Track collision point, PS projection...).
 template <unsigned DIM = 2>
 struct KDTreeBox {
   std::array<float, DIM> dimmin, dimmax;
 
   template <typename... Ts>
   KDTreeBox(Ts... dimargs) {
     static_assert(sizeof...(dimargs) == 2 * DIM, "Constructor requires 2*DIM args");
     std::vector<float> dims = {dimargs...};
     for (unsigned i = 0; i < DIM; ++i) {
       dimmin[i] = dims[2 * i];
       dimmax[i] = dims[2 * i + 1];
     }
   }
 
   KDTreeBox() {}
 };
 
 // Data stored in each KDTree node.
 // The dim1/dim2 fields are usually the duplication of some PFRecHit values
 // (eta/phi or x/y). But in some situations, phi field is shifted by +-2.Pi
 template <typename DATA, unsigned DIM = 2>
 struct KDTreeNodeInfo {
   DATA data;
   std::array<float, DIM> dims;
 
 public:
   KDTreeNodeInfo() {}
 
   template <typename... Ts>
   KDTreeNodeInfo(const DATA &d, Ts... dimargs) : data(d), dims{{dimargs...}} {}
   template <typename... Ts>
   bool operator>(const KDTreeNodeInfo &rhs) const {
     return (data > rhs.data);
   }
 };
 
 template <typename DATA, unsigned DIM = 2>
 struct KDTreeNodes {
   std::array<std::vector<float>, DIM> dims;
   std::vector<int> right;
   std::vector<DATA> data;
 
   int poolSize;
   int poolPos;
 
   constexpr KDTreeNodes() : poolSize(-1), poolPos(-1) {}
 
   bool empty() const { return poolPos == -1; }
   int size() const { return poolPos + 1; }
 
   void clear() {
     for (auto &dim : dims) {
       dim.clear();
       dim.shrink_to_fit();
     }
     right.clear();
     right.shrink_to_fit();
     data.clear();
     data.shrink_to_fit();
     poolSize = -1;
     poolPos = -1;
   }
 
   int getNextNode() {
     ++poolPos;
     return poolPos;
   }
 
   void build(int sizeData) {
     poolSize = sizeData * 2 - 1;
     for (auto &dim : dims) {
       dim.resize(poolSize);
     }
     right.resize(poolSize);
     data.resize(poolSize);
   };
 
   constexpr bool isLeaf(int right) const {
     // Valid values of right are always >= 2
     // index 0 is the root, and 1 is the first left node
     // Exploit index values 0 and 1 to mark which of dim1/dim2 is the
     // current one in recSearch() at the depth of the leaf.
     return right < 2;
   }
 
   bool isLeafIndex(int index) const { return isLeaf(right[index]); }
 };
 
 // Class that implements the KDTree partition of 2D space and
 // a closest point search algorithme.
 
 template <typename DATA, unsigned int DIM = 2>
 class KDTreeLinkerAlgo {
 public:
   // Dtor calls clear()
   ~KDTreeLinkerAlgo() { clear(); }
 
   // Here we build the KD tree from the "eltList" in the space define by "region".
   void build(std::vector<KDTreeNodeInfo<DATA, DIM> > &eltList, const KDTreeBox<DIM> &region);
 
   // Here we search in the KDTree for all points that would be
   // contained in the given searchbox. The founded points are stored in resRecHitList.
   void search(const KDTreeBox<DIM> &searchBox, std::vector<DATA> &resRecHitList);
 
   // This reurns true if the tree is empty
   bool empty() { return nodePool_.empty(); }
 
   // This returns the number of nodes + leaves in the tree
   // (nElements should be (size() +1)/2)
   int size() { return nodePool_.size(); }
 
   // This method clears all allocated structures.
   void clear() { clearTree(); }
 
 private:
   // The node pool allow us to do just 1 call to new for each tree building.
   KDTreeNodes<DATA, DIM> nodePool_;
 
   std::vector<DATA> *closestNeighbour;
   std::vector<KDTreeNodeInfo<DATA, DIM> > *initialEltList;
 
   //Fast median search with Wirth algorithm in eltList between low and high indexes.
   int medianSearch(int low, int high, int treeDepth) const;
 
   // Recursif kdtree builder. Is called by build()
   int recBuild(int low, int hight, int depth);
 
   // Recursif kdtree search. Is called by search()
   void recSearch(int current, const KDTreeBox<DIM> &trackBox, int depth = 0) const;
 
   // This method frees the KDTree.
   void clearTree() { nodePool_.clear(); }
 };
 
 //Implementation
 
 template <typename DATA, unsigned int DIM>
 void KDTreeLinkerAlgo<DATA, DIM>::build(std::vector<KDTreeNodeInfo<DATA, DIM> > &eltList,
                                         const KDTreeBox<DIM> &region) {
   if (!eltList.empty()) {
     initialEltList = &eltList;
 
     size_t size = initialEltList->size();
     nodePool_.build(size);
 
     // Here we build the KDTree
     int root = recBuild(0, size, 0);
     assert(root == 0);
 
     initialEltList = nullptr;
   }
 }
 
 //Fast median search with Wirth algorithm in eltList between low and high indexes.
 template <typename DATA, unsigned int DIM>
 int KDTreeLinkerAlgo<DATA, DIM>::medianSearch(int low, int high, int treeDepth) const {
   int nbrElts = high - low;
   int median = (nbrElts & 1) ? nbrElts / 2 : nbrElts / 2 - 1;
   median += low;
 
   int l = low;
   int m = high - 1;
 
   while (l < m) {
     KDTreeNodeInfo<DATA, DIM> elt = (*initialEltList)[median];
     int i = l;
     int j = m;
 
     do {
       // The even depth is associated to dim1 dimension
       // The odd one to dim2 dimension
       const unsigned thedim = treeDepth % DIM;
       while ((*initialEltList)[i].dims[thedim] < elt.dims[thedim])
         ++i;
       while ((*initialEltList)[j].dims[thedim] > elt.dims[thedim])
         --j;
 
       if (i <= j) {
         std::swap((*initialEltList)[i], (*initialEltList)[j]);
         i++;
         j--;
       }
     } while (i <= j);
     if (j < median)
       l = i;
     if (i > median)
       m = j;
   }
 
   return median;
 }
 
 template <typename DATA, unsigned int DIM>
 void KDTreeLinkerAlgo<DATA, DIM>::search(const KDTreeBox<DIM> &trackBox, std::vector<DATA> &recHits) {
   if (!empty()) {
     closestNeighbour = &recHits;
     recSearch(0, trackBox, 0);
     closestNeighbour = nullptr;
   }
 }
 
 template <typename DATA, unsigned int DIM>
 void KDTreeLinkerAlgo<DATA, DIM>::recSearch(int current, const KDTreeBox<DIM> &trackBox, int depth) const {
   // Iterate until leaf is found, or there are no children in the
   // search window. If search has to proceed on both children, proceed
   // the search to left child via recursion. Swap search window
   // dimension on alternate levels.
   while (true) {
     const int dimIndex = depth % DIM;
     int right = nodePool_.right[current];
     if (nodePool_.isLeaf(right)) {
       // If point inside the rectangle/area
       // Use intentionally bit-wise & instead of logical && for better
       // performance. It is faster to always do all comparisons than to
       // allow use of branches to not do some if any of the first ones
       // is false.
       bool isInside = true;
       for (unsigned i = 0; i < DIM; ++i) {
         float dimCurr = nodePool_.dims[i][current];
         isInside &= reco::branchless_and(dimCurr >= trackBox.dimmin[i], dimCurr <= trackBox.dimmax[i]);
       }
       if (isInside) {
         closestNeighbour->push_back(nodePool_.data[current]);
       }
       break;
     } else {
       float median = nodePool_.dims[dimIndex][current];
 
       bool goLeft = (trackBox.dimmin[dimIndex] <= median);
       bool goRight = (trackBox.dimmax[dimIndex] >= median);
 
       ++depth;
       if (goLeft & goRight) {
         int left = current + 1;
         recSearch(left, trackBox, depth);
         // continue with right
         current = right;
       } else if (goLeft) {
         ++current;
       } else if (goRight) {
         current = right;
       } else {
         break;
       }
     }
   }
 }
 
 template <typename DATA, unsigned int DIM>
 int KDTreeLinkerAlgo<DATA, DIM>::recBuild(int low, int high, int depth) {
   int portionSize = high - low;
 
   if (portionSize == 1) {  // Leaf case
     int leaf = nodePool_.getNextNode();
     const KDTreeNodeInfo<DATA, DIM> &info = (*initialEltList)[low];
     nodePool_.right[leaf] = 0;
     for (unsigned i = 0; i < DIM; ++i) {
       nodePool_.dims[i][leaf] = info.dims[i];
     }
     nodePool_.data[leaf] = info.data;
     return leaf;
 
   } else {  // Node case
 
     // The even depth is associated to dim1 dimension
     // The odd one to dim2 dimension
     int medianId = medianSearch(low, high, depth);
     int dimIndex = depth % DIM;
     float medianVal = (*initialEltList)[medianId].dims[dimIndex];
 
     // We create the node
     int nodeInd = nodePool_.getNextNode();
 
     ++depth;
     ++medianId;
 
     // We recursively build the son nodes
     int left = recBuild(low, medianId, depth);
     assert(nodeInd + 1 == left);
     int right = recBuild(medianId, high, depth);
     nodePool_.right[nodeInd] = right;
 
     nodePool_.dims[dimIndex][nodeInd] = medianVal;
 
     return nodeInd;
   }
 }
 
 #endif

This page was automatically generated by the 2.2.1 LXR engine. The LXR team