File indexing completed on 2025-04-13 22:49:35
0001 #ifndef KDTreeLinkerAlgoTemplated_h
0002 #define KDTreeLinkerAlgoTemplated_h
0003
0004 #include "DataFormats/Math/interface/logic.h"
0005
0006 #include <cassert>
0007 #include <vector>
0008 #include <array>
0009 #include <algorithm>
0010
0011
0012
0013
0014
0015 template <unsigned DIM = 2>
0016 struct KDTreeBox {
0017 std::array<float, DIM> dimmin, dimmax;
0018
0019 template <typename... Ts>
0020 KDTreeBox(Ts... dimargs) {
0021 static_assert(sizeof...(dimargs) == 2 * DIM, "Constructor requires 2*DIM args");
0022 std::vector<float> dims = {dimargs...};
0023 for (unsigned i = 0; i < DIM; ++i) {
0024 dimmin[i] = dims[2 * i];
0025 dimmax[i] = dims[2 * i + 1];
0026 }
0027 }
0028
0029 KDTreeBox() {}
0030 };
0031
0032
0033
0034
0035 template <typename DATA, unsigned DIM = 2>
0036 struct KDTreeNodeInfo {
0037 DATA data;
0038 std::array<float, DIM> dims;
0039
0040 public:
0041 KDTreeNodeInfo() {}
0042
0043 template <typename... Ts>
0044 KDTreeNodeInfo(const DATA &d, Ts... dimargs) : data(d), dims{{dimargs...}} {}
0045 template <typename... Ts>
0046 bool operator>(const KDTreeNodeInfo &rhs) const {
0047 return (data > rhs.data);
0048 }
0049 };
0050
0051 template <typename DATA, unsigned DIM = 2>
0052 struct KDTreeNodes {
0053 std::array<std::vector<float>, DIM> dims;
0054 std::vector<int> right;
0055 std::vector<DATA> data;
0056
0057 int poolSize;
0058 int poolPos;
0059
0060 constexpr KDTreeNodes() : poolSize(-1), poolPos(-1) {}
0061
0062 bool empty() const { return poolPos == -1; }
0063 int size() const { return poolPos + 1; }
0064
0065 void clear() {
0066 for (auto &dim : dims) {
0067 dim.clear();
0068 dim.shrink_to_fit();
0069 }
0070 right.clear();
0071 right.shrink_to_fit();
0072 data.clear();
0073 data.shrink_to_fit();
0074 poolSize = -1;
0075 poolPos = -1;
0076 }
0077
0078 int getNextNode() {
0079 ++poolPos;
0080 return poolPos;
0081 }
0082
0083 void build(int sizeData) {
0084 poolSize = sizeData * 2 - 1;
0085 for (auto &dim : dims) {
0086 dim.resize(poolSize);
0087 }
0088 right.resize(poolSize);
0089 data.resize(poolSize);
0090 };
0091
0092 constexpr bool isLeaf(int right) const {
0093
0094
0095
0096
0097 return right < 2;
0098 }
0099
0100 bool isLeafIndex(int index) const { return isLeaf(right[index]); }
0101 };
0102
0103
0104
0105
0106 template <typename DATA, unsigned int DIM = 2>
0107 class KDTreeLinkerAlgo {
0108 public:
0109
0110 ~KDTreeLinkerAlgo() { clear(); }
0111
0112
0113 void build(std::vector<KDTreeNodeInfo<DATA, DIM> > &eltList, const KDTreeBox<DIM> ®ion);
0114
0115
0116
0117 void search(const KDTreeBox<DIM> &searchBox, std::vector<DATA> &resRecHitList);
0118
0119
0120 bool empty() { return nodePool_.empty(); }
0121
0122
0123
0124 int size() { return nodePool_.size(); }
0125
0126
0127 void clear() { clearTree(); }
0128
0129 private:
0130
0131 KDTreeNodes<DATA, DIM> nodePool_;
0132
0133 std::vector<DATA> *closestNeighbour;
0134 std::vector<KDTreeNodeInfo<DATA, DIM> > *initialEltList;
0135
0136
0137 int medianSearch(int low, int high, int treeDepth) const;
0138
0139
0140 int recBuild(int low, int hight, int depth);
0141
0142
0143 void recSearch(int current, const KDTreeBox<DIM> &trackBox, int depth = 0) const;
0144
0145
0146 void clearTree() { nodePool_.clear(); }
0147 };
0148
0149
0150
0151 template <typename DATA, unsigned int DIM>
0152 void KDTreeLinkerAlgo<DATA, DIM>::build(std::vector<KDTreeNodeInfo<DATA, DIM> > &eltList,
0153 const KDTreeBox<DIM> ®ion) {
0154 if (!eltList.empty()) {
0155 initialEltList = &eltList;
0156
0157 size_t size = initialEltList->size();
0158 nodePool_.build(size);
0159
0160
0161 int root = recBuild(0, size, 0);
0162 assert(root == 0);
0163
0164 initialEltList = nullptr;
0165 }
0166 }
0167
0168
0169 template <typename DATA, unsigned int DIM>
0170 int KDTreeLinkerAlgo<DATA, DIM>::medianSearch(int low, int high, int treeDepth) const {
0171 int nbrElts = high - low;
0172 int median = (nbrElts & 1) ? nbrElts / 2 : nbrElts / 2 - 1;
0173 median += low;
0174
0175 int l = low;
0176 int m = high - 1;
0177
0178 while (l < m) {
0179 KDTreeNodeInfo<DATA, DIM> elt = (*initialEltList)[median];
0180 int i = l;
0181 int j = m;
0182
0183 do {
0184
0185
0186 const unsigned thedim = treeDepth % DIM;
0187 while ((*initialEltList)[i].dims[thedim] < elt.dims[thedim])
0188 ++i;
0189 while ((*initialEltList)[j].dims[thedim] > elt.dims[thedim])
0190 --j;
0191
0192 if (i <= j) {
0193 std::swap((*initialEltList)[i], (*initialEltList)[j]);
0194 i++;
0195 j--;
0196 }
0197 } while (i <= j);
0198 if (j < median)
0199 l = i;
0200 if (i > median)
0201 m = j;
0202 }
0203
0204 return median;
0205 }
0206
0207 template <typename DATA, unsigned int DIM>
0208 void KDTreeLinkerAlgo<DATA, DIM>::search(const KDTreeBox<DIM> &trackBox, std::vector<DATA> &recHits) {
0209 if (!empty()) {
0210 closestNeighbour = &recHits;
0211 recSearch(0, trackBox, 0);
0212 closestNeighbour = nullptr;
0213 }
0214 }
0215
0216 template <typename DATA, unsigned int DIM>
0217 void KDTreeLinkerAlgo<DATA, DIM>::recSearch(int current, const KDTreeBox<DIM> &trackBox, int depth) const {
0218
0219
0220
0221
0222 while (true) {
0223 const int dimIndex = depth % DIM;
0224 int right = nodePool_.right[current];
0225 if (nodePool_.isLeaf(right)) {
0226
0227
0228
0229
0230
0231 bool isInside = true;
0232 for (unsigned i = 0; i < DIM; ++i) {
0233 float dimCurr = nodePool_.dims[i][current];
0234 isInside &= reco::branchless_and(dimCurr >= trackBox.dimmin[i], dimCurr <= trackBox.dimmax[i]);
0235 }
0236 if (isInside) {
0237 closestNeighbour->push_back(nodePool_.data[current]);
0238 }
0239 break;
0240 } else {
0241 float median = nodePool_.dims[dimIndex][current];
0242
0243 bool goLeft = (trackBox.dimmin[dimIndex] <= median);
0244 bool goRight = (trackBox.dimmax[dimIndex] >= median);
0245
0246 ++depth;
0247 if (goLeft & goRight) {
0248 int left = current + 1;
0249 recSearch(left, trackBox, depth);
0250
0251 current = right;
0252 } else if (goLeft) {
0253 ++current;
0254 } else if (goRight) {
0255 current = right;
0256 } else {
0257 break;
0258 }
0259 }
0260 }
0261 }
0262
0263 template <typename DATA, unsigned int DIM>
0264 int KDTreeLinkerAlgo<DATA, DIM>::recBuild(int low, int high, int depth) {
0265 int portionSize = high - low;
0266
0267 if (portionSize == 1) {
0268 int leaf = nodePool_.getNextNode();
0269 const KDTreeNodeInfo<DATA, DIM> &info = (*initialEltList)[low];
0270 nodePool_.right[leaf] = 0;
0271 for (unsigned i = 0; i < DIM; ++i) {
0272 nodePool_.dims[i][leaf] = info.dims[i];
0273 }
0274 nodePool_.data[leaf] = info.data;
0275 return leaf;
0276
0277 } else {
0278
0279
0280
0281 int medianId = medianSearch(low, high, depth);
0282 int dimIndex = depth % DIM;
0283 float medianVal = (*initialEltList)[medianId].dims[dimIndex];
0284
0285
0286 int nodeInd = nodePool_.getNextNode();
0287
0288 ++depth;
0289 ++medianId;
0290
0291
0292 int left = recBuild(low, medianId, depth);
0293 assert(nodeInd + 1 == left);
0294 int right = recBuild(medianId, high, depth);
0295 nodePool_.right[nodeInd] = right;
0296
0297 nodePool_.dims[dimIndex][nodeInd] = medianVal;
0298
0299 return nodeInd;
0300 }
0301 }
0302
0303 #endif