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