Compadre  1.2.0
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Macros Pages
Compadre_PointCloudSearch.hpp
Go to the documentation of this file.
1 #ifndef _COMPADRE_POINTCLOUDSEARCH_HPP_
2 #define _COMPADRE_POINTCLOUDSEARCH_HPP_
3 
4 #include "Compadre_Typedefs.hpp"
6 #include "nanoflann.hpp"
7 #include <Kokkos_Core.hpp>
8 #include <memory>
9 
10 namespace Compadre {
11 
12 //class sort_indices
13 //{
14 // private:
15 // double* mparr;
16 // public:
17 // sort_indices(double* parr) : mparr(parr) {}
18 // bool operator()(int i, int j) const { return mparr[i]<mparr[j]; }
19 //};
20 
21 //! Custom RadiusResultSet for nanoflann that uses pre-allocated space for indices and radii
22 //! instead of using std::vec for std::pairs
23 template <typename _DistanceType, typename _IndexType = size_t>
25 
26  public:
27 
28  typedef _DistanceType DistanceType;
29  typedef _IndexType IndexType;
30 
36 
38  DistanceType radius_,
39  DistanceType* r_dist_, IndexType* i_dist_, const IndexType max_size_)
40  : radius(radius_), count(0), r_dist(r_dist_), i_dist(i_dist_), max_size(max_size_) {
41  init();
42  }
43 
44  void init() {}
45 
46  void clear() { count = 0; }
47 
48  size_t size() const { return count; }
49 
50  bool full() const { return true; }
51 
52  bool addPoint(DistanceType dist, IndexType index) {
53 
54  if (dist < radius) {
55  // would throw an exception here if count>=max_size, but this code is
56  // called inside of a parallel region so only an abort is possible,
57  // but this situation is recoverable
58  //
59  // instead, we increase count, but there is nowhere to store neighbors
60  // since we are working with pre-allocated space
61  // this will be discovered after returning from the search by comparing
62  // the count against the pre-allocate space dimensions
63  if (count<max_size) {
64  i_dist[count] = index;
65  r_dist[count] = dist;
66  }
67  count++;
68  }
69  return true;
70 
71  }
72 
73  DistanceType worstDist() const { return radius; }
74 
75  std::pair<IndexType, DistanceType> worst_item() const {
76  // just to verify this isn't ever called
77  compadre_kernel_assert_release(false && "worst_item() should not be called.");
78  }
79 
80  void sort() {
81  // puts closest neighbor as the first entry in the neighbor list
82  // leaves the rest unsorted
83 
84  if (count > 0) {
85 
86  // alternate sort for every entry, not currently being used
87  //int indices[count];
88  //for (int i=0; i<count; ++i) {
89  // indices[i] = i;
90  //}
91  //std::sort(indices, indices+count, sort_indices(r_dist));
92  //std::vector<IndexType> tmp_indices(count);
93  //std::vector<DistanceType> tmp_r(count);
94  //for (int i=0; i<count; ++i) {
95  // tmp_indices[i] = i_dist[indices[i]];
96  // tmp_r[i] = r_dist[indices[i]];
97  //}
98  //for (int i=0; i<count; ++i) {
99  // i_dist[i] = tmp_indices[i];
100  // r_dist[i] = tmp_r[i];
101  //}
102  IndexType loop_max = (count < max_size) ? count : max_size;
103 
104  DistanceType best_distance = std::numeric_limits<DistanceType>::max();
105  IndexType best_distance_index = 0;
106  int best_index = -1;
107  for (IndexType i=0; i<loop_max; ++i) {
108  if (r_dist[i] < best_distance) {
109  best_distance = r_dist[i];
110  best_distance_index = i_dist[i];
111  best_index = i;
112  }
113  }
114 
115  if (best_index != 0) {
116  auto tmp_ind = i_dist[0];
117  i_dist[0] = best_distance_index;
118  i_dist[best_index] = tmp_ind;
119  }
120  }
121  }
122 };
123 
124 
125 //! PointCloudSearch generates neighbor lists and window sizes for each target site
126 /*!
127 * Search methods can be run in dry-run mode, or not.
128 *
129 * #### When in dry-run mode:
130 *
131 * `neighbors_list` will be populated with number of neighbors found for each target site.
132 *
133 * This allows a user to know memory allocation needed before storage of neighbor indices.
134 *
135 * #### When not in dry-run mode:
136 *
137 * `neighbors_list_offsets` will be populated with offsets for values (dependent on method) determined by neighbor_list.
138 * If a 2D view for `neighbors_list` is used, then \f$ N(i,j+1) \f$ will store the \f$ j^{th} \f$ neighbor of \f$ i \f$,
139 * and \f$ N(i,0) \f$ will store the number of neighbors for target \f$ i \f$.
140 *
141 */
142 template <typename view_type>
144 
145  public:
146 
147  typedef nanoflann::KDTreeSingleIndexAdaptor<nanoflann::L2_Simple_Adaptor<double, PointCloudSearch<view_type> >,
149  typedef nanoflann::KDTreeSingleIndexAdaptor<nanoflann::L2_Simple_Adaptor<double, PointCloudSearch<view_type> >,
151  typedef nanoflann::KDTreeSingleIndexAdaptor<nanoflann::L2_Simple_Adaptor<double, PointCloudSearch<view_type> >,
153 
154  protected:
155 
156  //! source site coordinates
157  view_type _src_pts_view;
160 
161  std::shared_ptr<tree_type_1d> _tree_1d;
162  std::shared_ptr<tree_type_2d> _tree_2d;
163  std::shared_ptr<tree_type_3d> _tree_3d;
164 
165  public:
166 
167  PointCloudSearch(view_type src_pts_view, const local_index_type dimension = -1,
168  const local_index_type max_leaf = -1)
169  : _src_pts_view(src_pts_view),
170  _dim((dimension < 0) ? src_pts_view.extent(1) : dimension),
171  _max_leaf((max_leaf < 0) ? 10 : max_leaf) {
172  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename view_type::memory_space>::accessible==1)
173  && "Views passed to PointCloudSearch at construction should be accessible from the host.");
174  };
175 
177 
178  //! Returns a liberal estimated upper bound on number of neighbors to be returned by a neighbor search
179  //! for a given choice of dimension, basis size, and epsilon_multiplier. Assumes quasiuniform distribution
180  //! of points. This result can be used to size a preallocated neighbor_lists kokkos view.
181  static inline int getEstimatedNumberNeighborsUpperBound(int unisolvency_size, const int dimension, const double epsilon_multiplier) {
182  int multiplier = 1;
183  if (dimension==1) multiplier = 2;
184  return multiplier * 2.0 * unisolvency_size * pow(epsilon_multiplier, dimension) + 1; // +1 is for the number of neighbors entry needed in the first entry of each row
185  }
186 
187  //! Bounding box query method required by Nanoflann.
188  template <class BBOX> bool kdtree_get_bbox(BBOX& bb) const {return false;}
189 
190  //! Returns the number of source sites
191  inline int kdtree_get_point_count() const {return _src_pts_view.extent(0);}
192 
193  //! Returns the coordinate value of a point
194  inline double kdtree_get_pt(const int idx, int dim) const {return _src_pts_view(idx,dim);}
195 
196  //! Returns the distance between a point and a source site, given its index
197  inline double kdtree_distance(const double* queryPt, const int idx, long long sz) const {
198 
199  double distance = 0;
200  for (int i=0; i<_dim; ++i) {
201  distance += (_src_pts_view(idx,i)-queryPt[i])*(_src_pts_view(idx,i)-queryPt[i]);
202  }
203  return std::sqrt(distance);
204 
205  }
206 
207  void generateKDTree() {
208  if (_dim==1) {
209  _tree_1d = std::make_shared<tree_type_1d>(1, *this, nanoflann::KDTreeSingleIndexAdaptorParams(_max_leaf));
210  _tree_1d->buildIndex();
211  } else if (_dim==2) {
212  _tree_2d = std::make_shared<tree_type_2d>(2, *this, nanoflann::KDTreeSingleIndexAdaptorParams(_max_leaf));
213  _tree_2d->buildIndex();
214  } else if (_dim==3) {
215  _tree_3d = std::make_shared<tree_type_3d>(3, *this, nanoflann::KDTreeSingleIndexAdaptorParams(_max_leaf));
216  _tree_3d->buildIndex();
217  }
218  }
219 
220  /*! \brief Generates neighbor lists of 2D view by performing a radius search
221  where the radius to be searched is in the epsilons view.
222  If uniform_radius is given, then this overrides the epsilons view radii sizes.
223  Accepts 2D neighbor_lists without number_of_neighbors_list.
224  \param is_dry_run [in] - whether to do a dry-run (find neighbors, but don't store)
225  \param trg_pts_view [in] - target coordinates from which to seek neighbors
226  \param neighbor_lists [out] - 2D view of neighbor lists to be populated from search
227  \param epsilons [in/out] - radius to search, overwritten if uniform_radius != 0
228  \param uniform_radius [in] - double != 0 determines whether to overwrite all epsilons for uniform search
229  \param max_search_radius [in] - largest valid search (useful only for MPI jobs if halo size exists)
230  */
231  template <typename trg_view_type, typename neighbor_lists_view_type, typename epsilons_view_type>
232  size_t generate2DNeighborListsFromRadiusSearch(bool is_dry_run, trg_view_type trg_pts_view,
233  neighbor_lists_view_type neighbor_lists, epsilons_view_type epsilons,
234  const double uniform_radius = 0.0, double max_search_radius = 0.0) {
235 
236  // function does not populate epsilons, they must be prepopulated
237 
238  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename trg_view_type::memory_space>::accessible==1) &&
239  "Target coordinates view passed to generate2DNeighborListsFromRadiusSearch should be accessible from the host.");
240  compadre_assert_release((((int)trg_pts_view.extent(1))>=_dim) &&
241  "Target coordinates view passed to generate2DNeighborListsFromRadiusSearch must have \
242  second dimension as large as _dim.");
243  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename neighbor_lists_view_type::memory_space>::accessible==1) &&
244  "Views passed to generate2DNeighborListsFromRadiusSearch should be accessible from the host.");
245  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename epsilons_view_type::memory_space>::accessible==1) &&
246  "Views passed to generate2DNeighborListsFromRadiusSearch should be accessible from the host.");
247 
248  // loop size
249  const int num_target_sites = trg_pts_view.extent(0);
250 
251  if ((!_tree_1d && _dim==1) || (!_tree_2d && _dim==2) || (!_tree_3d && _dim==3)) {
252  this->generateKDTree();
253  }
254 
255  // check neighbor lists and epsilons view sizes
256  compadre_assert_release((neighbor_lists.extent(0)==(size_t)num_target_sites
257  && neighbor_lists.extent(1)>=1)
258  && "neighbor lists View does not have large enough dimensions");
259  compadre_assert_release((neighbor_lists_view_type::rank==2) && "neighbor_lists must be a 2D Kokkos view.");
260 
261  compadre_assert_release((epsilons.extent(0)==(size_t)num_target_sites)
262  && "epsilons View does not have the correct dimension");
263 
264  typedef Kokkos::View<double*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
265  scratch_double_view;
266 
267  typedef Kokkos::View<size_t*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
268  scratch_int_view;
269 
270  // determine scratch space size needed
271  int team_scratch_size = 0;
272  team_scratch_size += scratch_double_view::shmem_size(neighbor_lists.extent(1)); // distances
273  team_scratch_size += scratch_int_view::shmem_size(neighbor_lists.extent(1)); // indices
274  team_scratch_size += scratch_double_view::shmem_size(_dim); // target coordinate
275 
276  // maximum number of neighbors found over all target sites' neighborhoods
277  size_t max_num_neighbors = 0;
278  // part 2. do radius search using window size from knn search
279  // each row of neighbor lists is a neighbor list for the target site corresponding to that row
280  Kokkos::parallel_reduce("radius search", host_team_policy(num_target_sites, Kokkos::AUTO)
281  .set_scratch_size(0 /*shared memory level*/, Kokkos::PerTeam(team_scratch_size)),
282  KOKKOS_LAMBDA(const host_member_type& teamMember, size_t& t_max_num_neighbors) {
283 
284  // make unmanaged scratch views
285  scratch_double_view neighbor_distances(teamMember.team_scratch(0 /*shared memory*/), neighbor_lists.extent(1));
286  scratch_int_view neighbor_indices(teamMember.team_scratch(0 /*shared memory*/), neighbor_lists.extent(1));
287  scratch_double_view this_target_coord(teamMember.team_scratch(0 /*shared memory*/), _dim);
288 
289  size_t neighbors_found = 0;
290 
291  const int i = teamMember.league_rank();
292 
293  // set epsilons if radius is specified
294  if (uniform_radius > 0) epsilons(i) = uniform_radius;
295 
296  // needs furthest neighbor's distance for next portion
297  compadre_kernel_assert_release((epsilons(i)<=max_search_radius || max_search_radius==0) && "max_search_radius given (generally derived from the size of a halo region), and search radius needed would exceed this max_search_radius.");
298 
299  Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, neighbor_lists.extent(1)), [&](const int j) {
300  neighbor_indices(j) = 0;
301  neighbor_distances(j) = -1.0;
302  });
303  teamMember.team_barrier();
304 
305  Kokkos::single(Kokkos::PerTeam(teamMember), [&] () {
306  // target_coords is LayoutLeft on device and its HostMirror, so giving a pointer to
307  // this data would lead to a wrong result if the device is a GPU
308 
309  for (int j=0; j<_dim; ++j) {
310  this_target_coord(j) = trg_pts_view(i,j);
311  }
312 
313  nanoflann::SearchParams sp; // default parameters
314  Compadre::RadiusResultSet<double> rrs(epsilons(i)*epsilons(i), neighbor_distances.data(), neighbor_indices.data(), neighbor_lists.extent(1));
315  if (_dim==1) {
316  neighbors_found = _tree_1d->template radiusSearchCustomCallback<Compadre::RadiusResultSet<double> >(this_target_coord.data(), rrs, sp) ;
317  } else if (_dim==2) {
318  neighbors_found = _tree_2d->template radiusSearchCustomCallback<Compadre::RadiusResultSet<double> >(this_target_coord.data(), rrs, sp) ;
319  } else if (_dim==3) {
320  neighbors_found = _tree_3d->template radiusSearchCustomCallback<Compadre::RadiusResultSet<double> >(this_target_coord.data(), rrs, sp) ;
321  }
322 
323  t_max_num_neighbors = (neighbors_found > t_max_num_neighbors) ? neighbors_found : t_max_num_neighbors;
324 
325  // the number of neighbors is stored in column zero of the neighbor lists 2D array
326  neighbor_lists(i,0) = neighbors_found;
327 
328  // epsilons already scaled and then set by search radius
329  });
330  teamMember.team_barrier();
331 
332  // loop_bound so that we don't write into memory we don't have allocated
333  int loop_bound = (neighbors_found < neighbor_lists.extent(1)-1) ? neighbors_found : neighbor_lists.extent(1)-1;
334  // loop over each neighbor index and fill with a value
335  if (!is_dry_run) {
336  Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, loop_bound), [&](const int j) {
337  // cast to an whatever data type the 2D array of neighbor lists is using
338  neighbor_lists(i,j+1) = static_cast<typename std::remove_pointer<typename std::remove_pointer<typename neighbor_lists_view_type::data_type>::type>::type>(neighbor_indices(j));
339  });
340  teamMember.team_barrier();
341  }
342 
343  }, Kokkos::Max<size_t>(max_num_neighbors) );
344  Kokkos::fence();
345 
346  // check if max_num_neighbors will fit onto pre-allocated space
347  compadre_assert_release((neighbor_lists.extent(1) >= (max_num_neighbors+1) || is_dry_run)
348  && "neighbor_lists does not contain enough columns for the maximum number of neighbors needing to be stored.");
349 
350  return max_num_neighbors;
351  }
352 
353  /*! \brief Generates compressed row neighbor lists by performing a radius search
354  where the radius to be searched is in the epsilons view.
355  If uniform_radius is given, then this overrides the epsilons view radii sizes.
356  Accepts 1D neighbor_lists with 1D number_of_neighbors_list.
357  \param is_dry_run [in] - whether to do a dry-run (find neighbors, but don't store)
358  \param trg_pts_view [in] - target coordinates from which to seek neighbors
359  \param neighbor_lists [out] - 1D view of neighbor lists to be populated from search
360  \param number_of_neighbors_list [in/out] - number of neighbors for each target site
361  \param epsilons [in/out] - radius to search, overwritten if uniform_radius != 0
362  \param uniform_radius [in] - double != 0 determines whether to overwrite all epsilons for uniform search
363  \param max_search_radius [in] - largest valid search (useful only for MPI jobs if halo size exists)
364  */
365  template <typename trg_view_type, typename neighbor_lists_view_type, typename epsilons_view_type>
366  size_t generateCRNeighborListsFromRadiusSearch(bool is_dry_run, trg_view_type trg_pts_view,
367  neighbor_lists_view_type neighbor_lists, neighbor_lists_view_type number_of_neighbors_list,
368  epsilons_view_type epsilons, const double uniform_radius = 0.0, double max_search_radius = 0.0) {
369 
370  // function does not populate epsilons, they must be prepopulated
371 
372  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename trg_view_type::memory_space>::accessible==1) &&
373  "Target coordinates view passed to generateCRNeighborListsFromRadiusSearch should be accessible from the host.");
374  compadre_assert_release((((int)trg_pts_view.extent(1))>=_dim) &&
375  "Target coordinates view passed to generateCRNeighborListsFromRadiusSearch must have \
376  second dimension as large as _dim.");
377  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename neighbor_lists_view_type::memory_space>::accessible==1) &&
378  "Views passed to generateCRNeighborListsFromRadiusSearch should be accessible from the host.");
379  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename epsilons_view_type::memory_space>::accessible==1) &&
380  "Views passed to generateCRNeighborListsFromRadiusSearch should be accessible from the host.");
381 
382  // loop size
383  const int num_target_sites = trg_pts_view.extent(0);
384 
385  if ((!_tree_1d && _dim==1) || (!_tree_2d && _dim==2) || (!_tree_3d && _dim==3)) {
386  this->generateKDTree();
387  }
388 
389  compadre_assert_release((number_of_neighbors_list.extent(0)==(size_t)num_target_sites)
390  && "number_of_neighbors_list or neighbor lists View does not have large enough dimensions");
391  compadre_assert_release((neighbor_lists_view_type::rank==1) && "neighbor_lists must be a 1D Kokkos view.");
392 
393  neighbor_lists_view_type row_offsets;
394  int max_neighbor_list_row_storage_size = 1;
395  if (!is_dry_run) {
396  auto nla(CreateNeighborLists(neighbor_lists, number_of_neighbors_list));
397  max_neighbor_list_row_storage_size = nla.getMaxNumNeighbors();
398  Kokkos::resize(row_offsets, num_target_sites);
399  Kokkos::parallel_for(Kokkos::RangePolicy<host_execution_space>(0,num_target_sites), [&](const int i) {
400  row_offsets(i) = nla.getRowOffsetHost(i);
401  });
402  Kokkos::fence();
403  }
404 
405  compadre_assert_release((epsilons.extent(0)==(size_t)num_target_sites)
406  && "epsilons View does not have the correct dimension");
407 
408  typedef Kokkos::View<double*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
409  scratch_double_view;
410 
411  typedef Kokkos::View<size_t*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
412  scratch_int_view;
413 
414  // determine scratch space size needed
415  int team_scratch_size = 0;
416  team_scratch_size += scratch_double_view::shmem_size(max_neighbor_list_row_storage_size); // distances
417  team_scratch_size += scratch_int_view::shmem_size(max_neighbor_list_row_storage_size); // indices
418  team_scratch_size += scratch_double_view::shmem_size(_dim); // target coordinate
419 
420  // part 2. do radius search using window size from knn search
421  // each row of neighbor lists is a neighbor list for the target site corresponding to that row
422  Kokkos::parallel_for("radius search", host_team_policy(num_target_sites, Kokkos::AUTO)
423  .set_scratch_size(0 /*shared memory level*/, Kokkos::PerTeam(team_scratch_size)),
424  KOKKOS_LAMBDA(const host_member_type& teamMember) {
425 
426  // make unmanaged scratch views
427  scratch_double_view neighbor_distances(teamMember.team_scratch(0 /*shared memory*/), max_neighbor_list_row_storage_size);
428  scratch_int_view neighbor_indices(teamMember.team_scratch(0 /*shared memory*/), max_neighbor_list_row_storage_size);
429  scratch_double_view this_target_coord(teamMember.team_scratch(0 /*shared memory*/), _dim);
430 
431  size_t neighbors_found = 0;
432 
433  const int i = teamMember.league_rank();
434 
435  // set epsilons if radius is specified
436  if (uniform_radius > 0) epsilons(i) = uniform_radius;
437 
438  // needs furthest neighbor's distance for next portion
439  compadre_kernel_assert_release((epsilons(i)<=max_search_radius || max_search_radius==0) && "max_search_radius given (generally derived from the size of a halo region), and search radius needed would exceed this max_search_radius.");
440 
441  Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, max_neighbor_list_row_storage_size), [&](const int j) {
442  neighbor_indices(j) = 0;
443  neighbor_distances(j) = -1.0;
444  });
445  teamMember.team_barrier();
446 
447  Kokkos::single(Kokkos::PerTeam(teamMember), [&] () {
448  // target_coords is LayoutLeft on device and its HostMirror, so giving a pointer to
449  // this data would lead to a wrong result if the device is a GPU
450 
451  for (int j=0; j<_dim; ++j) {
452  this_target_coord(j) = trg_pts_view(i,j);
453  }
454 
455  nanoflann::SearchParams sp; // default parameters
456  Compadre::RadiusResultSet<double> rrs(epsilons(i)*epsilons(i), neighbor_distances.data(), neighbor_indices.data(), max_neighbor_list_row_storage_size);
457  if (_dim==1) {
458  neighbors_found = _tree_1d->template radiusSearchCustomCallback<Compadre::RadiusResultSet<double> >(this_target_coord.data(), rrs, sp) ;
459  } else if (_dim==2) {
460  neighbors_found = _tree_2d->template radiusSearchCustomCallback<Compadre::RadiusResultSet<double> >(this_target_coord.data(), rrs, sp) ;
461  } else if (_dim==3) {
462  neighbors_found = _tree_3d->template radiusSearchCustomCallback<Compadre::RadiusResultSet<double> >(this_target_coord.data(), rrs, sp) ;
463  }
464 
465  // the number of neighbors is stored in column zero of the neighbor lists 2D array
466  if (is_dry_run) {
467  number_of_neighbors_list(i) = neighbors_found;
468  } else {
469  compadre_kernel_assert_debug((neighbors_found==(size_t)number_of_neighbors_list(i))
470  && "Number of neighbors found changed since dry-run.");
471  }
472 
473  // epsilons already scaled and then set by search radius
474  });
475  teamMember.team_barrier();
476 
477  // loop_bound so that we don't write into memory we don't have allocated
478  int loop_bound = neighbors_found;
479  // loop over each neighbor index and fill with a value
480  if (!is_dry_run) {
481  Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, loop_bound), [&](const int j) {
482  // cast to an whatever data type the 2D array of neighbor lists is using
483  neighbor_lists(row_offsets(i)+j) = static_cast<typename std::remove_pointer<typename std::remove_pointer<typename neighbor_lists_view_type::data_type>::type>::type>(neighbor_indices(j));
484  });
485  teamMember.team_barrier();
486  }
487 
488  });
489  Kokkos::fence();
490  if (is_dry_run) {
491  int total_storage_size = 0;
492  Kokkos::parallel_reduce("total number of neighbors over all lists", Kokkos::RangePolicy<host_execution_space>(0, number_of_neighbors_list.extent(0)),
493  KOKKOS_LAMBDA(const int i, int& t_total_num_neighbors) {
494  t_total_num_neighbors += number_of_neighbors_list(i);
495  }, Kokkos::Sum<int>(total_storage_size));
496  Kokkos::fence();
497  return (size_t)(total_storage_size);
498  } else {
499  return (size_t)(number_of_neighbors_list(num_target_sites-1)+row_offsets(num_target_sites-1));
500  }
501  }
502 
503 
504  /*! \brief Generates neighbor lists as 2D view by performing a k-nearest neighbor search
505  Only accepts 2D neighbor_lists without number_of_neighbors_list.
506  \param is_dry_run [in] - whether to do a dry-run (find neighbors, but don't store)
507  \param trg_pts_view [in] - target coordinates from which to seek neighbors
508  \param neighbor_lists [out] - 2D view of neighbor lists to be populated from search
509  \param epsilons [in/out] - radius to search, overwritten if uniform_radius != 0
510  \param neighbors_needed [in] - k neighbors needed as a minimum
511  \param epsilon_multiplier [in] - distance to kth neighbor multiplied by epsilon_multiplier for follow-on radius search
512  \param max_search_radius [in] - largest valid search (useful only for MPI jobs if halo size exists)
513  */
514  template <typename trg_view_type, typename neighbor_lists_view_type, typename epsilons_view_type>
515  size_t generate2DNeighborListsFromKNNSearch(bool is_dry_run, trg_view_type trg_pts_view,
516  neighbor_lists_view_type neighbor_lists, epsilons_view_type epsilons,
517  const int neighbors_needed, const double epsilon_multiplier = 1.6,
518  double max_search_radius = 0.0) {
519 
520  // First, do a knn search (removes need for guessing initial search radius)
521 
522  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename trg_view_type::memory_space>::accessible==1) &&
523  "Target coordinates view passed to generate2DNeighborListsFromKNNSearch should be accessible from the host.");
524  compadre_assert_release((((int)trg_pts_view.extent(1))>=_dim) &&
525  "Target coordinates view passed to generate2DNeighborListsFromRadiusSearch must have \
526  second dimension as large as _dim.");
527  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename neighbor_lists_view_type::memory_space>::accessible==1) &&
528  "Views passed to generate2DNeighborListsFromKNNSearch should be accessible from the host.");
529  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename epsilons_view_type::memory_space>::accessible==1) &&
530  "Views passed to generate2DNeighborListsFromKNNSearch should be accessible from the host.");
531 
532  // loop size
533  const int num_target_sites = trg_pts_view.extent(0);
534 
535  if ((!_tree_1d && _dim==1) || (!_tree_2d && _dim==2) || (!_tree_3d && _dim==3)) {
536  this->generateKDTree();
537  }
538  Kokkos::fence();
539 
540  compadre_assert_release((num_target_sites==0 || // sizes don't matter when there are no targets
541  (neighbor_lists.extent(0)==(size_t)num_target_sites
542  && neighbor_lists.extent(1)>=(size_t)(neighbors_needed+1)))
543  && "neighbor lists View does not have large enough dimensions");
544  compadre_assert_release((neighbor_lists_view_type::rank==2) && "neighbor_lists must be a 2D Kokkos view.");
545 
546  compadre_assert_release((epsilons.extent(0)==(size_t)num_target_sites)
547  && "epsilons View does not have the correct dimension");
548 
549  typedef Kokkos::View<double*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
550  scratch_double_view;
551 
552  typedef Kokkos::View<size_t*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
553  scratch_int_view;
554 
555  // determine scratch space size needed
556  int team_scratch_size = 0;
557  team_scratch_size += scratch_double_view::shmem_size(neighbor_lists.extent(1)); // distances
558  team_scratch_size += scratch_int_view::shmem_size(neighbor_lists.extent(1)); // indices
559  team_scratch_size += scratch_double_view::shmem_size(_dim); // target coordinate
560 
561  // minimum number of neighbors found over all target sites' neighborhoods
562  size_t min_num_neighbors = 0;
563  //
564  // part 1. do knn search for neighbors needed for unisolvency
565  // each row of neighbor lists is a neighbor list for the target site corresponding to that row
566  //
567  // as long as neighbor_lists can hold the number of neighbors_needed, we don't need to check
568  // that the maximum number of neighbors will fit into neighbor_lists
569  //
570  Kokkos::parallel_reduce("knn search", host_team_policy(num_target_sites, Kokkos::AUTO)
571  .set_scratch_size(0 /*shared memory level*/, Kokkos::PerTeam(team_scratch_size)),
572  KOKKOS_LAMBDA(const host_member_type& teamMember, size_t& t_min_num_neighbors) {
573 
574  // make unmanaged scratch views
575  scratch_double_view neighbor_distances(teamMember.team_scratch(0 /*shared memory*/), neighbor_lists.extent(1));
576  scratch_int_view neighbor_indices(teamMember.team_scratch(0 /*shared memory*/), neighbor_lists.extent(1));
577  scratch_double_view this_target_coord(teamMember.team_scratch(0 /*shared memory*/), _dim);
578 
579  size_t neighbors_found = 0;
580 
581  const int i = teamMember.league_rank();
582 
583  Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, neighbor_lists.extent(1)), [=](const int j) {
584  neighbor_indices(j) = 0;
585  neighbor_distances(j) = -1.0;
586  });
587 
588  teamMember.team_barrier();
589  Kokkos::single(Kokkos::PerTeam(teamMember), [&] () {
590  // target_coords is LayoutLeft on device and its HostMirror, so giving a pointer to
591  // this data would lead to a wrong result if the device is a GPU
592 
593  for (int j=0; j<_dim; ++j) {
594  this_target_coord(j) = trg_pts_view(i,j);
595  }
596 
597  if (_dim==1) {
598  neighbors_found = _tree_1d->knnSearch(this_target_coord.data(), neighbors_needed,
599  neighbor_indices.data(), neighbor_distances.data()) ;
600  } else if (_dim==2) {
601  neighbors_found = _tree_2d->knnSearch(this_target_coord.data(), neighbors_needed,
602  neighbor_indices.data(), neighbor_distances.data()) ;
603  } else if (_dim==3) {
604  neighbors_found = _tree_3d->knnSearch(this_target_coord.data(), neighbors_needed,
605  neighbor_indices.data(), neighbor_distances.data()) ;
606  }
607 
608  // get minimum number of neighbors found over all target sites' neighborhoods
609  t_min_num_neighbors = (neighbors_found < t_min_num_neighbors) ? neighbors_found : t_min_num_neighbors;
610 
611  // scale by epsilon_multiplier to window from location where the last neighbor was found
612  epsilons(i) = (neighbor_distances(neighbors_found-1) > 0) ?
613  std::sqrt(neighbor_distances(neighbors_found-1))*(epsilon_multiplier+1e-14) : 1e-14*epsilon_multiplier;
614  // the only time the second case using 1e-14 is used is when either zero neighbors or exactly one
615  // neighbor (neighbor is target site) is found. when the follow on radius search is conducted, the one
616  // neighbor (target site) will not be found if left at 0, so any positive amount will do, however 1e-14
617  // should is small enough to ensure that other neighbors are not found
618 
619  // needs furthest neighbor's distance for next portion
620  compadre_kernel_assert_release((neighbors_found<neighbor_lists.extent(1) || is_dry_run)
621  && "Neighbors found exceeded storage capacity in neighbor list.");
622  compadre_kernel_assert_release((epsilons(i)<=max_search_radius || max_search_radius==0 || is_dry_run)
623  && "max_search_radius given (generally derived from the size of a halo region), \
624  and search radius needed would exceed this max_search_radius.");
625  // neighbor_distances stores squared distances from neighbor to target, as returned by nanoflann
626  });
627  }, Kokkos::Min<size_t>(min_num_neighbors) );
628  Kokkos::fence();
629 
630  // if no target sites, then min_num_neighbors is set to neighbors_needed
631  // which also avoids min_num_neighbors being improperly set by min reduction
632  if (num_target_sites==0) min_num_neighbors = neighbors_needed;
633 
634  // Next, check that we found the neighbors_needed number that we require for unisolvency
635  compadre_assert_release((num_target_sites==0 || (min_num_neighbors>=(size_t)neighbors_needed))
636  && "Neighbor search failed to find number of neighbors needed for unisolvency.");
637 
638  // call a radius search using values now stored in epsilons
639  size_t max_num_neighbors = generate2DNeighborListsFromRadiusSearch(is_dry_run, trg_pts_view, neighbor_lists,
640  epsilons, 0.0 /*don't set uniform radius*/, max_search_radius);
641 
642  return max_num_neighbors;
643  }
644 
645  /*! \brief Generates compressed row neighbor lists by performing a k-nearest neighbor search
646  Only accepts 1D neighbor_lists with 1D number_of_neighbors_list.
647  \param is_dry_run [in] - whether to do a dry-run (find neighbors, but don't store)
648  \param trg_pts_view [in] - target coordinates from which to seek neighbors
649  \param neighbor_lists [out] - 1D view of neighbor lists to be populated from search
650  \param number_of_neighbors_list [in/out] - number of neighbors for each target site
651  \param epsilons [in/out] - radius to search, overwritten if uniform_radius != 0
652  \param neighbors_needed [in] - k neighbors needed as a minimum
653  \param epsilon_multiplier [in] - distance to kth neighbor multiplied by epsilon_multiplier for follow-on radius search
654  \param max_search_radius [in] - largest valid search (useful only for MPI jobs if halo size exists)
655  */
656  template <typename trg_view_type, typename neighbor_lists_view_type, typename epsilons_view_type>
657  size_t generateCRNeighborListsFromKNNSearch(bool is_dry_run, trg_view_type trg_pts_view,
658  neighbor_lists_view_type neighbor_lists, neighbor_lists_view_type number_of_neighbors_list,
659  epsilons_view_type epsilons, const int neighbors_needed, const double epsilon_multiplier = 1.6,
660  double max_search_radius = 0.0) {
661 
662  // First, do a knn search (removes need for guessing initial search radius)
663 
664  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename trg_view_type::memory_space>::accessible==1) &&
665  "Target coordinates view passed to generateCRNeighborListsFromKNNSearch should be accessible from the host.");
666  compadre_assert_release((((int)trg_pts_view.extent(1))>=_dim) &&
667  "Target coordinates view passed to generateCRNeighborListsFromRadiusSearch must have \
668  second dimension as large as _dim.");
669  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename neighbor_lists_view_type::memory_space>::accessible==1) &&
670  "Views passed to generateCRNeighborListsFromKNNSearch should be accessible from the host.");
671  compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename epsilons_view_type::memory_space>::accessible==1) &&
672  "Views passed to generateCRNeighborListsFromKNNSearch should be accessible from the host.");
673 
674  // loop size
675  const int num_target_sites = trg_pts_view.extent(0);
676 
677  if ((!_tree_1d && _dim==1) || (!_tree_2d && _dim==2) || (!_tree_3d && _dim==3)) {
678  this->generateKDTree();
679  }
680  Kokkos::fence();
681 
682  compadre_assert_release((number_of_neighbors_list.extent(0)==(size_t)num_target_sites )
683  && "number_of_neighbors_list or neighbor lists View does not have large enough dimensions");
684  compadre_assert_release((neighbor_lists_view_type::rank==1) && "neighbor_lists must be a 1D Kokkos view.");
685 
686  int last_row_offset = 0;
687  // if dry-run, neighbors_needed, else max over previous dry-run
688  int max_neighbor_list_row_storage_size = neighbors_needed;
689  if (!is_dry_run) {
690  auto nla(CreateNeighborLists(neighbor_lists, number_of_neighbors_list));
691  max_neighbor_list_row_storage_size = nla.getMaxNumNeighbors();
692  last_row_offset = nla.getRowOffsetHost(num_target_sites-1);
693  }
694 
695  compadre_assert_release((epsilons.extent(0)==(size_t)num_target_sites)
696  && "epsilons View does not have the correct dimension");
697 
698  typedef Kokkos::View<double*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
699  scratch_double_view;
700 
701  typedef Kokkos::View<size_t*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
702  scratch_int_view;
703 
704  // determine scratch space size needed
705  int team_scratch_size = 0;
706  team_scratch_size += scratch_double_view::shmem_size(max_neighbor_list_row_storage_size); // distances
707  team_scratch_size += scratch_int_view::shmem_size(max_neighbor_list_row_storage_size); // indices
708  team_scratch_size += scratch_double_view::shmem_size(_dim); // target coordinate
709 
710  // minimum number of neighbors found over all target sites' neighborhoods
711  size_t min_num_neighbors = 0;
712  //
713  // part 1. do knn search for neighbors needed for unisolvency
714  // each row of neighbor lists is a neighbor list for the target site corresponding to that row
715  //
716  // as long as neighbor_lists can hold the number of neighbors_needed, we don't need to check
717  // that the maximum number of neighbors will fit into neighbor_lists
718  //
719  Kokkos::parallel_reduce("knn search", host_team_policy(num_target_sites, Kokkos::AUTO)
720  .set_scratch_size(0 /*shared memory level*/, Kokkos::PerTeam(team_scratch_size)),
721  KOKKOS_LAMBDA(const host_member_type& teamMember, size_t& t_min_num_neighbors) {
722 
723  // make unmanaged scratch views
724  scratch_double_view neighbor_distances(teamMember.team_scratch(0 /*shared memory*/), max_neighbor_list_row_storage_size);
725  scratch_int_view neighbor_indices(teamMember.team_scratch(0 /*shared memory*/), max_neighbor_list_row_storage_size);
726  scratch_double_view this_target_coord(teamMember.team_scratch(0 /*shared memory*/), _dim);
727 
728  size_t neighbors_found = 0;
729 
730  const int i = teamMember.league_rank();
731 
732  Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, max_neighbor_list_row_storage_size), [=](const int j) {
733  neighbor_indices(j) = 0;
734  neighbor_distances(j) = -1.0;
735  });
736 
737  teamMember.team_barrier();
738  Kokkos::single(Kokkos::PerTeam(teamMember), [&] () {
739  // target_coords is LayoutLeft on device and its HostMirror, so giving a pointer to
740  // this data would lead to a wrong result if the device is a GPU
741 
742  for (int j=0; j<_dim; ++j) {
743  this_target_coord(j) = trg_pts_view(i,j);
744  }
745 
746  if (_dim==1) {
747  neighbors_found = _tree_1d->knnSearch(this_target_coord.data(), neighbors_needed,
748  neighbor_indices.data(), neighbor_distances.data()) ;
749  } else if (_dim==2) {
750  neighbors_found = _tree_2d->knnSearch(this_target_coord.data(), neighbors_needed,
751  neighbor_indices.data(), neighbor_distances.data()) ;
752  } else if (_dim==3) {
753  neighbors_found = _tree_3d->knnSearch(this_target_coord.data(), neighbors_needed,
754  neighbor_indices.data(), neighbor_distances.data()) ;
755  }
756 
757  // get minimum number of neighbors found over all target sites' neighborhoods
758  t_min_num_neighbors = (neighbors_found < t_min_num_neighbors) ? neighbors_found : t_min_num_neighbors;
759 
760  // scale by epsilon_multiplier to window from location where the last neighbor was found
761  epsilons(i) = (neighbor_distances(neighbors_found-1) > 0) ?
762  std::sqrt(neighbor_distances(neighbors_found-1))*(epsilon_multiplier+1e-14) : 1e-14*epsilon_multiplier;
763  // the only time the second case using 1e-14 is used is when either zero neighbors or exactly one
764  // neighbor (neighbor is target site) is found. when the follow on radius search is conducted, the one
765  // neighbor (target site) will not be found if left at 0, so any positive amount will do, however 1e-14
766  // should is small enough to ensure that other neighbors are not found
767 
768  compadre_kernel_assert_release((epsilons(i)<=max_search_radius || max_search_radius==0 || is_dry_run)
769  && "max_search_radius given (generally derived from the size of a halo region), \
770  and search radius needed would exceed this max_search_radius.");
771  // neighbor_distances stores squared distances from neighbor to target, as returned by nanoflann
772  });
773  }, Kokkos::Min<size_t>(min_num_neighbors) );
774  Kokkos::fence();
775 
776  // if no target sites, then min_num_neighbors is set to neighbors_needed
777  // which also avoids min_num_neighbors being improperly set by min reduction
778  if (num_target_sites==0) min_num_neighbors = neighbors_needed;
779 
780  // Next, check that we found the neighbors_needed number that we require for unisolvency
781  compadre_assert_release((num_target_sites==0 || (min_num_neighbors>=(size_t)neighbors_needed))
782  && "Neighbor search failed to find number of neighbors needed for unisolvency.");
783 
784  // call a radius search using values now stored in epsilons
785  generateCRNeighborListsFromRadiusSearch(is_dry_run, trg_pts_view, neighbor_lists,
786  number_of_neighbors_list, epsilons, 0.0 /*don't set uniform radius*/, max_search_radius);
787 
788  if (is_dry_run) {
789  int total_storage_size = 0;
790  Kokkos::parallel_reduce("total number of neighbors over all lists", Kokkos::RangePolicy<host_execution_space>(0, number_of_neighbors_list.extent(0)),
791  KOKKOS_LAMBDA(const int i, int& t_total_num_neighbors) {
792  t_total_num_neighbors += number_of_neighbors_list(i);
793  }, Kokkos::Sum<int>(total_storage_size));
794  Kokkos::fence();
795  return (size_t)(total_storage_size);
796  } else {
797  return (size_t)(number_of_neighbors_list(num_target_sites-1)+last_row_offset);
798  }
799  }
800 }; // PointCloudSearch
801 
802 //! CreatePointCloudSearch allows for the construction of an object of type PointCloudSearch with template deduction
803 template <typename view_type>
804 PointCloudSearch<view_type> CreatePointCloudSearch(view_type src_view, const local_index_type dimensions = -1, const local_index_type max_leaf = -1) {
805  return PointCloudSearch<view_type>(src_view, dimensions, max_leaf);
806 }
807 
808 } // Compadre
809 
810 #endif
size_t generate2DNeighborListsFromKNNSearch(bool is_dry_run, trg_view_type trg_pts_view, neighbor_lists_view_type neighbor_lists, epsilons_view_type epsilons, const int neighbors_needed, const double epsilon_multiplier=1.6, double max_search_radius=0.0)
Generates neighbor lists as 2D view by performing a k-nearest neighbor search Only accepts 2D neighbo...
nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor< double, PointCloudSearch< view_type > >, PointCloudSearch< view_type >, 3 > tree_type_3d
import subprocess import os import re import math import sys import argparse d d num_target_sites
#define compadre_kernel_assert_release(condition)
compadre_kernel_assert_release is similar to compadre_assert_release, but is a call on the device...
int kdtree_get_point_count() const
Returns the number of source sites.
PointCloudSearch< view_type > CreatePointCloudSearch(view_type src_view, const local_index_type dimensions=-1, const local_index_type max_leaf=-1)
CreatePointCloudSearch allows for the construction of an object of type PointCloudSearch with templat...
RadiusResultSet(DistanceType radius_, DistanceType *r_dist_, IndexType *i_dist_, const IndexType max_size_)
size_t generateCRNeighborListsFromKNNSearch(bool is_dry_run, trg_view_type trg_pts_view, neighbor_lists_view_type neighbor_lists, neighbor_lists_view_type number_of_neighbors_list, epsilons_view_type epsilons, const int neighbors_needed, const double epsilon_multiplier=1.6, double max_search_radius=0.0)
Generates compressed row neighbor lists by performing a k-nearest neighbor search Only accepts 1D nei...
bool kdtree_get_bbox(BBOX &bb) const
Bounding box query method required by Nanoflann.
Kokkos::TeamPolicy< host_execution_space > host_team_policy
size_t generate2DNeighborListsFromRadiusSearch(bool is_dry_run, trg_view_type trg_pts_view, neighbor_lists_view_type neighbor_lists, epsilons_view_type epsilons, const double uniform_radius=0.0, double max_search_radius=0.0)
Generates neighbor lists of 2D view by performing a radius search where the radius to be searched is ...
size_t generateCRNeighborListsFromRadiusSearch(bool is_dry_run, trg_view_type trg_pts_view, neighbor_lists_view_type neighbor_lists, neighbor_lists_view_type number_of_neighbors_list, epsilons_view_type epsilons, const double uniform_radius=0.0, double max_search_radius=0.0)
Generates compressed row neighbor lists by performing a radius search where the radius to be searched...
std::pair< IndexType, DistanceType > worst_item() const
NeighborLists< view_type > CreateNeighborLists(view_type neighbor_lists, view_type number_of_neighbors_list)
CreateNeighborLists allows for the construction of an object of type NeighborLists with template dedu...
Custom RadiusResultSet for nanoflann that uses pre-allocated space for indices and radii instead of u...
PointCloudSearch(view_type src_pts_view, const local_index_type dimension=-1, const local_index_type max_leaf=-1)
nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor< double, PointCloudSearch< view_type > >, PointCloudSearch< view_type >, 2 > tree_type_2d
bool addPoint(DistanceType dist, IndexType index)
std::shared_ptr< tree_type_3d > _tree_3d
std::shared_ptr< tree_type_2d > _tree_2d
double kdtree_distance(const double *queryPt, const int idx, long long sz) const
Returns the distance between a point and a source site, given its index.
static int getEstimatedNumberNeighborsUpperBound(int unisolvency_size, const int dimension, const double epsilon_multiplier)
Returns a liberal estimated upper bound on number of neighbors to be returned by a neighbor search fo...
PointCloudSearch generates neighbor lists and window sizes for each target site.
host_team_policy::member_type host_member_type
int local_index_type
std::shared_ptr< tree_type_1d > _tree_1d
nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor< double, PointCloudSearch< view_type > >, PointCloudSearch< view_type >, 1 > tree_type_1d
double kdtree_get_pt(const int idx, int dim) const
Returns the coordinate value of a point.
view_type _src_pts_view
source site coordinates
#define compadre_assert_release(condition)
compadre_assert_release is used for assertions that should always be checked, but generally are not e...
#define compadre_kernel_assert_debug(condition)