GetFEM  5.4.2
bgeot_kdtree.cc
1 /*===========================================================================
2 
3  Copyright (C) 2004-2020 Julien Pommier
4 
5  This file is a part of GetFEM
6 
7  GetFEM is free software; you can redistribute it and/or modify it
8  under the terms of the GNU Lesser General Public License as published
9  by the Free Software Foundation; either version 3 of the License, or
10  (at your option) any later version along with the GCC Runtime Library
11  Exception either version 3.1 or (at your option) any later version.
12  This program is distributed in the hope that it will be useful, but
13  WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
14  or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
15  License and GCC Runtime Library Exception for more details.
16  You should have received a copy of the GNU Lesser General Public License
17  along with this program; if not, write to the Free Software Foundation,
18  Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.
19 
20 ===========================================================================*/
21 
22 
23 #include "getfem/bgeot_kdtree.h"
24 
25 namespace bgeot {
26 
27  /* leafs contains a list of points */
28  struct kdtree_leaf : public kdtree_elt_base {
29  kdtree_tab_type::const_iterator it;
30  kdtree_leaf(kdtree_tab_type::const_iterator begin,
31  kdtree_tab_type::const_iterator end) :
32  kdtree_elt_base(unsigned(std::distance(begin,end))) { it = begin; }
33  };
34 
35  struct kdtree_node : public kdtree_elt_base {
36  scalar_type split_v;
37  std::unique_ptr<kdtree_elt_base> left, right; /* left: <=v, right: >v */
38  kdtree_node(scalar_type v, std::unique_ptr<kdtree_elt_base> &&left_,
39  std::unique_ptr<kdtree_elt_base> &&right_) :
40  kdtree_elt_base(0), split_v(v),
41  left(std::move(left_)), right(std::move(right_)) {}
42  };
43 
44  typedef kdtree_tab_type::iterator ITER;
45 
46  /* sorting of kdtree_tab_type with respect to a component of the points */
47  struct component_sort {
48  unsigned dir;
49  component_sort(unsigned d) : dir(d) {}
50  bool operator()(const index_node_pair& a, const index_node_pair& b)
51  { return a.n.at(dir) < b.n.at(dir); }
52  };
53  /* splitting of kdtree_tab_type with respect to a given value */
54  /*struct component_split {
55  unsigned dir; scalar_type v;
56  component_split(unsigned d, scalar_type v_) : dir(d), v(v_) {}
57  bool operator()(const index_node_pair& a)
58  { return (a.n.at(dir) <= v); }
59  };
60  */
61 
62  static ITER partition(ITER begin, ITER end, unsigned dir, scalar_type median) {
63  --end;
64  do {
65  while (begin <= end && (*begin).n.at(dir) <= median) ++begin;
66  while (begin <= end && (*end).n.at(dir) > median) --end;
67  if (begin < end) { begin->swap(*end); ++begin; --end; } else break;
68  } while (1);
69  return begin;
70  }
71  /*
72  build (recursively) a complete tree for points in the interval
73  [begin, end[ dir is the splitting direction.
74  (may be ameliorated doing a single sort on each component at the begining)
75  */
76  static std::unique_ptr<kdtree_elt_base> build_tree_(ITER begin,
77  ITER end, unsigned dir) {
78  if (begin == end) return 0;
79  size_type npts = std::distance(begin,end);
80  if (npts > kdtree_elt_base::PTS_PER_LEAF) {
81  ITER itmedian;
82  scalar_type median;
83  size_type N = begin->n.size();
84  unsigned ndir_tests = unsigned(dir/N); dir = unsigned(dir % N);
85  if (npts > 50) {
86  /* too much points for an exact median: estimation of the median .. */
87  std::vector<index_node_pair> v(30);
88  //random_sample(begin,end,v.begin(),v.end());
89  for (size_type i=0; i < v.size(); ++i)
90  v[i] = begin[rand() % npts];
91  std::sort(v.begin(), v.end(), component_sort(dir));
92  median = (v[v.size()/2-1].n.at(dir)+v[v.size()/2].n.at(dir))/2;
93  //itmedian = std::partition(begin,end,component_split(dir, median));
94  itmedian = partition(begin,end,dir,median);
95  /*ITER it = begin; cout << "}"; while (it < end) {
96  if (it == itmedian) cout << "<|>"; else cout << " ";
97  cout << (*it).n[dir];
98  ++it;
99  } cout << "}\n";*/
100 
101  } else {
102  /* exact median computation */
103  std::sort(begin, end, component_sort(dir));
104  itmedian = begin + npts/2 - 1;
105  median = (*itmedian).n[dir];
106  while (itmedian < end && (*itmedian).n[dir] == median) itmedian++;
107  }
108  if (itmedian == end) /* could not split the set (all points have same value for component 'dir' !) */
109  if (ndir_tests == N-1) /* tested all N direction ? so all points are strictly the same */
110  return std::make_unique<kdtree_leaf>(begin,end);
111  else return std::make_unique<kdtree_node>
112  (median,
113  build_tree_(begin, itmedian, unsigned((dir+1)%N + (ndir_tests+1)*N)), std::unique_ptr<kdtree_node>());
114  else { /* the general case */
115  assert((*itmedian).n[dir] >= median && median >= (*(itmedian-1)).n[dir]);
116  return std::make_unique<kdtree_node>
117  (median, build_tree_(begin, itmedian, unsigned((dir+1)%N)),
118  build_tree_(itmedian,end, unsigned((dir+1)%N)));
119  }
120  } else {
121  return std::make_unique<kdtree_leaf>(begin,end);
122  }
123  }
124 
125  /* avoid pushing too much arguments on the stack for points_in_box_ */
126  struct points_in_box_data_ {
127  base_node::const_iterator bmin;
128  base_node::const_iterator bmax;
129  kdtree_tab_type *ipts;
130  size_type N;
131  };
132 
133  /* recursive lookup for points inside a given box */
134  static void points_in_box_(const points_in_box_data_& p,
135  const kdtree_elt_base *t, unsigned dir) {
136  if (!t->isleaf()) {
137  const kdtree_node *tn = static_cast<const kdtree_node*>(t);
138  if (p.bmin[dir] <= tn->split_v && tn->left.get())
139  points_in_box_(p, tn->left.get(), unsigned((dir+1)%p.N));
140  if (p.bmax[dir] > tn->split_v && tn->right)
141  points_in_box_(p, tn->right.get(), unsigned((dir+1)%p.N));
142  } else {
143  const kdtree_leaf *tl = static_cast<const kdtree_leaf*>(t);
144  kdtree_tab_type::const_iterator itpt = tl->it;
145  for (size_type i=tl->n;i; --i, ++itpt) {
146  bool is_in = true;
147  base_node::const_iterator it=itpt->n.const_begin();
148  for (size_type k=0; k < p.N; ++k) {
149  //cout << "test: k=" << k << ", " << it[k] << ", p.bmin[k]=" << p.bmin[k] << ", p.bmax[k]=" << p.bmax[k] << "\n";
150  if (it[k] < p.bmin[k] || it[k] > p.bmax[k]) {
151  is_in = false; break;
152  }
153  }
154  if (is_in) p.ipts->push_back(*itpt);
155  }
156  }
157  }
158 
159  /* avoid pushing too much arguments on the stack for nearest_neighbor_ */
160  struct nearest_neighbor_data_ {
161  base_node::const_iterator pos;
162  index_node_pair *ipt;
163  size_type N;
164  mutable scalar_type dist2;
165  base_node::iterator vec_to_tree_elm;
166  };
167 
168  /* recursive lookup for the nearest neighbor point at a given position */
169  static void nearest_neighbor_assist(const nearest_neighbor_data_& p,
170  const kdtree_elt_base *t, unsigned dir) {
171  scalar_type dist2(0);
172  for (size_type k=0; k < p.N; ++k)
173  dist2 += p.vec_to_tree_elm[k] * p.vec_to_tree_elm[k];
174  if (dist2 > p.dist2 && p.dist2 > scalar_type(0)) return;
175 
176  if (!t->isleaf()) {
177  const kdtree_node *tn = static_cast<const kdtree_node*>(t);
178  scalar_type tmp = p.vec_to_tree_elm[dir];
179  scalar_type dist = p.pos[dir] - tn->split_v;
180  if (tn->left.get()) {
181  if (dist > tmp) p.vec_to_tree_elm[dir] = dist;
182  nearest_neighbor_assist(p, tn->left.get(), unsigned((dir+1)%p.N));
183  p.vec_to_tree_elm[dir] = tmp;
184  }
185  if (tn->right) {
186  if (-dist > tmp) p.vec_to_tree_elm[dir] = -dist;
187  nearest_neighbor_assist(p, tn->right.get(), unsigned((dir+1)%p.N));
188  p.vec_to_tree_elm[dir] = tmp;
189  }
190  } else {
191  // find the nearest neighbor inside the leaf
192  const kdtree_leaf *tl = static_cast<const kdtree_leaf*>(t);
193  kdtree_tab_type::const_iterator itpt = tl->it;
194  for (size_type i=tl->n;i; --i, ++itpt) {
195  dist2 = scalar_type(0);
196  base_node::const_iterator it=itpt->n.const_begin();
197  for (size_type k=0; k < p.N; ++k) {
198  scalar_type dist = it[k] - p.pos[k];
199  dist2 += dist * dist;
200  }
201  if (dist2 < p.dist2 || p.dist2 < scalar_type(0)){
202  *(p.ipt) = *itpt;
203  p.dist2 = dist2;
204  }
205  }
206  }
207  }
208 
209  static void nearest_neighbor_main(const nearest_neighbor_data_& p,
210  const kdtree_elt_base *t, unsigned dir) {
211  if (!t->isleaf()) {
212  const kdtree_node *tn = static_cast<const kdtree_node*>(t);
213  scalar_type dist = p.pos[dir] - tn->split_v;
214  if ((dist <= scalar_type(0) && tn->left.get()) || !tn->right.get()) {
215  nearest_neighbor_main(p, tn->left.get(), unsigned((dir+1)%p.N));
216  } else if (tn->right.get()) {
217  nearest_neighbor_main(p, tn->right.get(), unsigned((dir+1)%p.N));
218  } else {
219  assert(false);
220  }
221  // check the possibility of points at the opposite side of the current
222  // tree node which are closer to pos as the current minimum distance
223  if (dist * dist <= p.dist2) {
224  for (size_type k=0; k < p.N; ++k) p.vec_to_tree_elm[k] = 0.;
225  if ((dist <= scalar_type(0) && tn->right.get()) || !tn->left.get()) {
226  p.vec_to_tree_elm[dir] = -dist;
227  nearest_neighbor_assist(p, tn->right.get(), unsigned((dir+1)%p.N));
228  } else if (tn->left.get()) {
229  p.vec_to_tree_elm[dir] = dist;
230  nearest_neighbor_assist(p, tn->left.get(), unsigned((dir+1)%p.N));
231  }
232  }
233  } else {
234  // find the nearest neighbor inside the leaf which contains pos
235  nearest_neighbor_assist(p, t, dir);
236  }
237  }
238 
239  void kdtree::clear_tree() { tree = std::unique_ptr<kdtree_elt_base>(); }
240 
241  void kdtree::points_in_box(kdtree_tab_type &ipts,
242  const base_node &min,
243  const base_node &max) {
244  ipts.resize(0);
245  if (tree == 0) { tree = build_tree_(pts.begin(), pts.end(), 0); if (!tree) return; }
246  base_node bmin(min), bmax(max);
247  for (size_type i=0; i < bmin.size(); ++i) if (bmin[i] > bmax[i]) return;
248  points_in_box_data_ p;
249  p.bmin = bmin.const_begin(); p.bmax = bmax.const_begin();
250  p.ipts = &ipts; p.N = N;
251  points_in_box_(p, tree.get(), 0);
252  }
253 
254  scalar_type kdtree::nearest_neighbor(index_node_pair &ipt,
255  const base_node &pos) {
256 
257  ipt.i = size_type(-1);
258  if (tree == 0) {
259  tree = build_tree_(pts.begin(), pts.end(), 0);
260  if (!tree) return scalar_type(-1);
261  }
262  nearest_neighbor_data_ p;
263  p.pos = pos.const_begin();
264  p.ipt = &ipt;
265  p.N = N;
266  p.dist2 = scalar_type(-1);
267  base_node tmp(N);
268  p.vec_to_tree_elm = tmp.begin();
269  nearest_neighbor_main(p, tree.get(), 0);
270  return p.dist2;
271  }
272 }
bgeot_kdtree.h
Simple implementation of a KD-tree.
bgeot::kdtree_tab_type
std::vector< index_node_pair > kdtree_tab_type
store a set of points with associated indexes.
Definition: bgeot_kdtree.h:68
bgeot::size_type
size_t size_type
used as the common size type in the library
Definition: bgeot_poly.h:49
bgeot
Basic Geometric Tools.
Definition: bgeot_convex_ref.cc:27