Compadre 1.5.5
Loading...
Searching...
No Matches
Compadre_PointCloudSearch.hpp
Go to the documentation of this file.
1#ifndef _COMPADRE_POINTCLOUDSEARCH_HPP_
2#define _COMPADRE_POINTCLOUDSEARCH_HPP_
3
6#include "nanoflann.hpp"
7#include <Kokkos_Core.hpp>
8#include <memory>
9
10namespace 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
23template <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*/
142template <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 * std::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
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 typedef Kokkos::View<global_index_type*, typename neighbor_lists_view_type::array_layout,
394 typename neighbor_lists_view_type::memory_space, typename neighbor_lists_view_type::memory_traits> row_offsets_view_type;
395 row_offsets_view_type row_offsets;
396 int max_neighbor_list_row_storage_size = 1;
397 if (!is_dry_run) {
398 auto nla = CreateNeighborLists(neighbor_lists, number_of_neighbors_list);
399 max_neighbor_list_row_storage_size = nla.getMaxNumNeighbors();
400 Kokkos::resize(row_offsets, num_target_sites);
401 Kokkos::fence();
402 Kokkos::parallel_for(Kokkos::RangePolicy<host_execution_space>(0,num_target_sites), [&](const int i) {
403 row_offsets(i) = nla.getRowOffsetHost(i);
404 });
405 Kokkos::fence();
406 }
407
408 compadre_assert_release((epsilons.extent(0)==(size_t)num_target_sites)
409 && "epsilons View does not have the correct dimension");
410
411 typedef Kokkos::View<double*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
412 scratch_double_view;
413
414 typedef Kokkos::View<size_t*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
415 scratch_int_view;
416
417 // determine scratch space size needed
418 int team_scratch_size = 0;
419 team_scratch_size += scratch_double_view::shmem_size(max_neighbor_list_row_storage_size); // distances
420 team_scratch_size += scratch_int_view::shmem_size(max_neighbor_list_row_storage_size); // indices
421 team_scratch_size += scratch_double_view::shmem_size(_dim); // target coordinate
422
423 // part 2. do radius search using window size from knn search
424 // each row of neighbor lists is a neighbor list for the target site corresponding to that row
425 Kokkos::parallel_for("radius search", host_team_policy(num_target_sites, Kokkos::AUTO)
426 .set_scratch_size(0 /*shared memory level*/, Kokkos::PerTeam(team_scratch_size)),
427 KOKKOS_LAMBDA(const host_member_type& teamMember) {
428
429 // make unmanaged scratch views
430 scratch_double_view neighbor_distances(teamMember.team_scratch(0 /*shared memory*/), max_neighbor_list_row_storage_size);
431 scratch_int_view neighbor_indices(teamMember.team_scratch(0 /*shared memory*/), max_neighbor_list_row_storage_size);
432 scratch_double_view this_target_coord(teamMember.team_scratch(0 /*shared memory*/), _dim);
433
434 size_t neighbors_found = 0;
435
436 const int i = teamMember.league_rank();
437
438 // set epsilons if radius is specified
439 if (uniform_radius > 0) epsilons(i) = uniform_radius;
440
441 // needs furthest neighbor's distance for next portion
442 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.");
443
444 Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, max_neighbor_list_row_storage_size), [&](const int j) {
445 neighbor_indices(j) = 0;
446 neighbor_distances(j) = -1.0;
447 });
448 teamMember.team_barrier();
449
450 Kokkos::single(Kokkos::PerTeam(teamMember), [&] () {
451 // target_coords is LayoutLeft on device and its HostMirror, so giving a pointer to
452 // this data would lead to a wrong result if the device is a GPU
453
454 for (int j=0; j<_dim; ++j) {
455 this_target_coord(j) = trg_pts_view(i,j);
456 }
457
458 nanoflann::SearchParams sp; // default parameters
459 Compadre::RadiusResultSet<double> rrs(epsilons(i)*epsilons(i), neighbor_distances.data(), neighbor_indices.data(), max_neighbor_list_row_storage_size);
460 if (_dim==1) {
461 neighbors_found = _tree_1d->template radiusSearchCustomCallback<Compadre::RadiusResultSet<double> >(this_target_coord.data(), rrs, sp) ;
462 } else if (_dim==2) {
463 neighbors_found = _tree_2d->template radiusSearchCustomCallback<Compadre::RadiusResultSet<double> >(this_target_coord.data(), rrs, sp) ;
464 } else if (_dim==3) {
465 neighbors_found = _tree_3d->template radiusSearchCustomCallback<Compadre::RadiusResultSet<double> >(this_target_coord.data(), rrs, sp) ;
466 }
467
468 // we check that neighbors found doesn't differ from dry-run or we store neighbors_found
469 // no check that neighbors found stay the same if uniform_radius specified (!=0)
470 if (is_dry_run || uniform_radius!=0.0) {
471 number_of_neighbors_list(i) = neighbors_found;
472 } else {
473 compadre_kernel_assert_debug((neighbors_found==(size_t)number_of_neighbors_list(i))
474 && "Number of neighbors found changed since dry-run.");
475 }
476
477 // epsilons already scaled and then set by search radius
478 });
479 teamMember.team_barrier();
480
481 // loop_bound so that we don't write into memory we don't have allocated
482 int loop_bound = neighbors_found;
483 // loop over each neighbor index and fill with a value
484 if (!is_dry_run) {
485 Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, loop_bound), [&](const int j) {
486 // cast to an whatever data type the 2D array of neighbor lists is using
487 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));
488 });
489 teamMember.team_barrier();
490 }
491
492 });
493 Kokkos::fence();
494 auto nla = CreateNeighborLists(number_of_neighbors_list);
495 return nla.getTotalNeighborsOverAllListsHost();
496 }
497
498
499 /*! \brief Generates neighbor lists as 2D view by performing a k-nearest neighbor search
500 Only accepts 2D neighbor_lists without number_of_neighbors_list.
501 \param is_dry_run [in] - whether to do a dry-run (find neighbors, but don't store)
502 \param trg_pts_view [in] - target coordinates from which to seek neighbors
503 \param neighbor_lists [out] - 2D view of neighbor lists to be populated from search
504 \param epsilons [in/out] - radius to search, overwritten if uniform_radius != 0
505 \param neighbors_needed [in] - k neighbors needed as a minimum
506 \param epsilon_multiplier [in] - distance to kth neighbor multiplied by epsilon_multiplier for follow-on radius search
507 \param max_search_radius [in] - largest valid search (useful only for MPI jobs if halo size exists)
508 */
509 template <typename trg_view_type, typename neighbor_lists_view_type, typename epsilons_view_type>
510 size_t generate2DNeighborListsFromKNNSearch(bool is_dry_run, trg_view_type trg_pts_view,
511 neighbor_lists_view_type neighbor_lists, epsilons_view_type epsilons,
512 const int neighbors_needed, const double epsilon_multiplier = 1.6,
513 double max_search_radius = 0.0) {
514
515 // First, do a knn search (removes need for guessing initial search radius)
516
517 compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename trg_view_type::memory_space>::accessible==1) &&
518 "Target coordinates view passed to generate2DNeighborListsFromKNNSearch should be accessible from the host.");
519 compadre_assert_release((((int)trg_pts_view.extent(1))>=_dim) &&
520 "Target coordinates view passed to generate2DNeighborListsFromRadiusSearch must have \
521 second dimension as large as _dim.");
522 compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename neighbor_lists_view_type::memory_space>::accessible==1) &&
523 "Views passed to generate2DNeighborListsFromKNNSearch should be accessible from the host.");
524 compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename epsilons_view_type::memory_space>::accessible==1) &&
525 "Views passed to generate2DNeighborListsFromKNNSearch should be accessible from the host.");
526
527 // loop size
528 const int num_target_sites = trg_pts_view.extent(0);
529
530 if ((!_tree_1d && _dim==1) || (!_tree_2d && _dim==2) || (!_tree_3d && _dim==3)) {
531 this->generateKDTree();
532 }
533 Kokkos::fence();
534
535 compadre_assert_release((num_target_sites==0 || // sizes don't matter when there are no targets
536 (neighbor_lists.extent(0)==(size_t)num_target_sites
537 && neighbor_lists.extent(1)>=(size_t)(neighbors_needed+1)))
538 && "neighbor lists View does not have large enough dimensions");
539 compadre_assert_release((neighbor_lists_view_type::rank==2) && "neighbor_lists must be a 2D Kokkos view.");
540
541 compadre_assert_release((epsilons.extent(0)==(size_t)num_target_sites)
542 && "epsilons View does not have the correct dimension");
543
544 typedef Kokkos::View<double*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
545 scratch_double_view;
546
547 typedef Kokkos::View<size_t*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
548 scratch_int_view;
549
550 // determine scratch space size needed
551 int team_scratch_size = 0;
552 team_scratch_size += scratch_double_view::shmem_size(neighbor_lists.extent(1)); // distances
553 team_scratch_size += scratch_int_view::shmem_size(neighbor_lists.extent(1)); // indices
554 team_scratch_size += scratch_double_view::shmem_size(_dim); // target coordinate
555
556 // minimum number of neighbors found over all target sites' neighborhoods
557 size_t min_num_neighbors = 0;
558 //
559 // part 1. do knn search for neighbors needed for unisolvency
560 // each row of neighbor lists is a neighbor list for the target site corresponding to that row
561 //
562 // as long as neighbor_lists can hold the number of neighbors_needed, we don't need to check
563 // that the maximum number of neighbors will fit into neighbor_lists
564 //
565 Kokkos::parallel_reduce("knn search", host_team_policy(num_target_sites, Kokkos::AUTO)
566 .set_scratch_size(0 /*shared memory level*/, Kokkos::PerTeam(team_scratch_size)),
567 KOKKOS_LAMBDA(const host_member_type& teamMember, size_t& t_min_num_neighbors) {
568
569 // make unmanaged scratch views
570 scratch_double_view neighbor_distances(teamMember.team_scratch(0 /*shared memory*/), neighbor_lists.extent(1));
571 scratch_int_view neighbor_indices(teamMember.team_scratch(0 /*shared memory*/), neighbor_lists.extent(1));
572 scratch_double_view this_target_coord(teamMember.team_scratch(0 /*shared memory*/), _dim);
573
574 size_t neighbors_found = 0;
575
576 const int i = teamMember.league_rank();
577
578 Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, neighbor_lists.extent(1)), [=](const int j) {
579 neighbor_indices(j) = 0;
580 neighbor_distances(j) = -1.0;
581 });
582
583 teamMember.team_barrier();
584 Kokkos::single(Kokkos::PerTeam(teamMember), [&] () {
585 // target_coords is LayoutLeft on device and its HostMirror, so giving a pointer to
586 // this data would lead to a wrong result if the device is a GPU
587
588 for (int j=0; j<_dim; ++j) {
589 this_target_coord(j) = trg_pts_view(i,j);
590 }
591
592 if (_dim==1) {
593 neighbors_found = _tree_1d->knnSearch(this_target_coord.data(), neighbors_needed,
594 neighbor_indices.data(), neighbor_distances.data()) ;
595 } else if (_dim==2) {
596 neighbors_found = _tree_2d->knnSearch(this_target_coord.data(), neighbors_needed,
597 neighbor_indices.data(), neighbor_distances.data()) ;
598 } else if (_dim==3) {
599 neighbors_found = _tree_3d->knnSearch(this_target_coord.data(), neighbors_needed,
600 neighbor_indices.data(), neighbor_distances.data()) ;
601 }
602
603 // get minimum number of neighbors found over all target sites' neighborhoods
604 t_min_num_neighbors = (neighbors_found < t_min_num_neighbors) ? neighbors_found : t_min_num_neighbors;
605
606 // scale by epsilon_multiplier to window from location where the last neighbor was found
607 epsilons(i) = (neighbor_distances(neighbors_found-1) > 0) ?
608 std::sqrt(neighbor_distances(neighbors_found-1))*epsilon_multiplier : 1e-14*epsilon_multiplier;
609 // the only time the second case using 1e-14 is used is when either zero neighbors or exactly one
610 // neighbor (neighbor is target site) is found. when the follow on radius search is conducted, the one
611 // neighbor (target site) will not be found if left at 0, so any positive amount will do, however 1e-14
612 // should is small enough to ensure that other neighbors are not found
613
614 // needs furthest neighbor's distance for next portion
615 compadre_kernel_assert_release((neighbors_found<neighbor_lists.extent(1) || is_dry_run)
616 && "Neighbors found exceeded storage capacity in neighbor list.");
617 compadre_kernel_assert_release((epsilons(i)<=max_search_radius || max_search_radius==0 || is_dry_run)
618 && "max_search_radius given (generally derived from the size of a halo region), \
619 and search radius needed would exceed this max_search_radius.");
620 // neighbor_distances stores squared distances from neighbor to target, as returned by nanoflann
621 });
622 }, Kokkos::Min<size_t>(min_num_neighbors) );
623 Kokkos::fence();
624
625 // if no target sites, then min_num_neighbors is set to neighbors_needed
626 // which also avoids min_num_neighbors being improperly set by min reduction
627 if (num_target_sites==0) min_num_neighbors = neighbors_needed;
628
629 // Next, check that we found the neighbors_needed number that we require for unisolvency
630 compadre_assert_release((num_target_sites==0 || (min_num_neighbors>=(size_t)neighbors_needed))
631 && "Neighbor search failed to find number of neighbors needed for unisolvency.");
632
633 // call a radius search using values now stored in epsilons
634 size_t max_num_neighbors = generate2DNeighborListsFromRadiusSearch(is_dry_run, trg_pts_view, neighbor_lists,
635 epsilons, 0.0 /*don't set uniform radius*/, max_search_radius);
636
637 return max_num_neighbors;
638 }
639
640 /*! \brief Generates compressed row neighbor lists by performing a k-nearest neighbor search
641 Only accepts 1D neighbor_lists with 1D number_of_neighbors_list.
642 \param is_dry_run [in] - whether to do a dry-run (find neighbors, but don't store)
643 \param trg_pts_view [in] - target coordinates from which to seek neighbors
644 \param neighbor_lists [out] - 1D view of neighbor lists to be populated from search
645 \param number_of_neighbors_list [in/out] - number of neighbors for each target site
646 \param epsilons [in/out] - radius to search, overwritten if uniform_radius != 0
647 \param neighbors_needed [in] - k neighbors needed as a minimum
648 \param epsilon_multiplier [in] - distance to kth neighbor multiplied by epsilon_multiplier for follow-on radius search
649 \param max_search_radius [in] - largest valid search (useful only for MPI jobs if halo size exists)
650 */
651 template <typename trg_view_type, typename neighbor_lists_view_type, typename epsilons_view_type>
652 size_t generateCRNeighborListsFromKNNSearch(bool is_dry_run, trg_view_type trg_pts_view,
653 neighbor_lists_view_type neighbor_lists, neighbor_lists_view_type number_of_neighbors_list,
654 epsilons_view_type epsilons, const int neighbors_needed, const double epsilon_multiplier = 1.6,
655 double max_search_radius = 0.0) {
656
657 // First, do a knn search (removes need for guessing initial search radius)
658
659 compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename trg_view_type::memory_space>::accessible==1) &&
660 "Target coordinates view passed to generateCRNeighborListsFromKNNSearch should be accessible from the host.");
661 compadre_assert_release((((int)trg_pts_view.extent(1))>=_dim) &&
662 "Target coordinates view passed to generateCRNeighborListsFromRadiusSearch must have \
663 second dimension as large as _dim.");
664 compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename neighbor_lists_view_type::memory_space>::accessible==1) &&
665 "Views passed to generateCRNeighborListsFromKNNSearch should be accessible from the host.");
666 compadre_assert_release((Kokkos::SpaceAccessibility<host_execution_space, typename epsilons_view_type::memory_space>::accessible==1) &&
667 "Views passed to generateCRNeighborListsFromKNNSearch should be accessible from the host.");
668
669 // loop size
670 const int num_target_sites = trg_pts_view.extent(0);
671
672 if ((!_tree_1d && _dim==1) || (!_tree_2d && _dim==2) || (!_tree_3d && _dim==3)) {
673 this->generateKDTree();
674 }
675 Kokkos::fence();
676
677 compadre_assert_release((number_of_neighbors_list.extent(0)==(size_t)num_target_sites )
678 && "number_of_neighbors_list or neighbor lists View does not have large enough dimensions");
679 compadre_assert_release((neighbor_lists_view_type::rank==1) && "neighbor_lists must be a 1D Kokkos view.");
680
681 // if dry-run, neighbors_needed, else max over previous dry-run
682 int max_neighbor_list_row_storage_size = neighbors_needed;
683 if (!is_dry_run) {
684 auto nla = CreateNeighborLists(neighbor_lists, number_of_neighbors_list);
685 max_neighbor_list_row_storage_size = nla.getMaxNumNeighbors();
686 }
687
688 compadre_assert_release((epsilons.extent(0)==(size_t)num_target_sites)
689 && "epsilons View does not have the correct dimension");
690
691 typedef Kokkos::View<double*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
692 scratch_double_view;
693
694 typedef Kokkos::View<size_t*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >
695 scratch_int_view;
696
697 // determine scratch space size needed
698 int team_scratch_size = 0;
699 team_scratch_size += scratch_double_view::shmem_size(max_neighbor_list_row_storage_size); // distances
700 team_scratch_size += scratch_int_view::shmem_size(max_neighbor_list_row_storage_size); // indices
701 team_scratch_size += scratch_double_view::shmem_size(_dim); // target coordinate
702
703 // minimum number of neighbors found over all target sites' neighborhoods
704 size_t min_num_neighbors = 0;
705 //
706 // part 1. do knn search for neighbors needed for unisolvency
707 // each row of neighbor lists is a neighbor list for the target site corresponding to that row
708 //
709 // as long as neighbor_lists can hold the number of neighbors_needed, we don't need to check
710 // that the maximum number of neighbors will fit into neighbor_lists
711 //
712 Kokkos::parallel_reduce("knn search", host_team_policy(num_target_sites, Kokkos::AUTO)
713 .set_scratch_size(0 /*shared memory level*/, Kokkos::PerTeam(team_scratch_size)),
714 KOKKOS_LAMBDA(const host_member_type& teamMember, size_t& t_min_num_neighbors) {
715
716 // make unmanaged scratch views
717 scratch_double_view neighbor_distances(teamMember.team_scratch(0 /*shared memory*/), max_neighbor_list_row_storage_size);
718 scratch_int_view neighbor_indices(teamMember.team_scratch(0 /*shared memory*/), max_neighbor_list_row_storage_size);
719 scratch_double_view this_target_coord(teamMember.team_scratch(0 /*shared memory*/), _dim);
720
721 size_t neighbors_found = 0;
722
723 const int i = teamMember.league_rank();
724
725 Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, max_neighbor_list_row_storage_size), [=](const int j) {
726 neighbor_indices(j) = 0;
727 neighbor_distances(j) = -1.0;
728 });
729
730 teamMember.team_barrier();
731 Kokkos::single(Kokkos::PerTeam(teamMember), [&] () {
732 // target_coords is LayoutLeft on device and its HostMirror, so giving a pointer to
733 // this data would lead to a wrong result if the device is a GPU
734
735 for (int j=0; j<_dim; ++j) {
736 this_target_coord(j) = trg_pts_view(i,j);
737 }
738
739 if (_dim==1) {
740 neighbors_found = _tree_1d->knnSearch(this_target_coord.data(), neighbors_needed,
741 neighbor_indices.data(), neighbor_distances.data()) ;
742 } else if (_dim==2) {
743 neighbors_found = _tree_2d->knnSearch(this_target_coord.data(), neighbors_needed,
744 neighbor_indices.data(), neighbor_distances.data()) ;
745 } else if (_dim==3) {
746 neighbors_found = _tree_3d->knnSearch(this_target_coord.data(), neighbors_needed,
747 neighbor_indices.data(), neighbor_distances.data()) ;
748 }
749
750 // get minimum number of neighbors found over all target sites' neighborhoods
751 t_min_num_neighbors = (neighbors_found < t_min_num_neighbors) ? neighbors_found : t_min_num_neighbors;
752
753 // scale by epsilon_multiplier to window from location where the last neighbor was found
754 epsilons(i) = (neighbor_distances(neighbors_found-1) > 0) ?
755 std::sqrt(neighbor_distances(neighbors_found-1))*epsilon_multiplier : 1e-14*epsilon_multiplier;
756 // the only time the second case using 1e-14 is used is when either zero neighbors or exactly one
757 // neighbor (neighbor is target site) is found. when the follow on radius search is conducted, the one
758 // neighbor (target site) will not be found if left at 0, so any positive amount will do, however 1e-14
759 // should is small enough to ensure that other neighbors are not found
760
761 compadre_kernel_assert_release((epsilons(i)<=max_search_radius || max_search_radius==0 || is_dry_run)
762 && "max_search_radius given (generally derived from the size of a halo region), \
763 and search radius needed would exceed this max_search_radius.");
764 // neighbor_distances stores squared distances from neighbor to target, as returned by nanoflann
765 });
766 }, Kokkos::Min<size_t>(min_num_neighbors) );
767 Kokkos::fence();
768
769 // if no target sites, then min_num_neighbors is set to neighbors_needed
770 // which also avoids min_num_neighbors being improperly set by min reduction
771 if (num_target_sites==0) min_num_neighbors = neighbors_needed;
772
773 // Next, check that we found the neighbors_needed number that we require for unisolvency
774 compadre_assert_release((num_target_sites==0 || (min_num_neighbors>=(size_t)neighbors_needed))
775 && "Neighbor search failed to find number of neighbors needed for unisolvency.");
776
777 // call a radius search using values now stored in epsilons
778 generateCRNeighborListsFromRadiusSearch(is_dry_run, trg_pts_view, neighbor_lists,
779 number_of_neighbors_list, epsilons, 0.0 /*don't set uniform radius*/, max_search_radius);
780
781 auto nla = CreateNeighborLists(number_of_neighbors_list);
782 return nla.getTotalNeighborsOverAllListsHost();
783 }
784}; // PointCloudSearch
785
786//! CreatePointCloudSearch allows for the construction of an object of type PointCloudSearch with template deduction
787template <typename view_type>
788PointCloudSearch<view_type> CreatePointCloudSearch(view_type src_view, const local_index_type dimensions = -1, const local_index_type max_leaf = -1) {
789 return PointCloudSearch<view_type>(src_view, dimensions, max_leaf);
790}
791
792} // Compadre
793
794#endif
int local_index_type
Kokkos::TeamPolicy< host_execution_space > host_team_policy
#define compadre_kernel_assert_debug(condition)
std::size_t global_index_type
#define compadre_kernel_assert_release(condition)
compadre_kernel_assert_release is similar to compadre_assert_release, but is a call on the device,...
host_team_policy::member_type host_member_type
#define compadre_assert_release(condition)
compadre_assert_release is used for assertions that should always be checked, but generally are not e...
PointCloudSearch generates neighbor lists and window sizes for each target site.
bool kdtree_get_bbox(BBOX &bb) const
Bounding box query method required by Nanoflann.
double kdtree_get_pt(const int idx, int dim) const
Returns the coordinate value of a point.
std::shared_ptr< tree_type_1d > _tree_1d
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...
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.
view_type _src_pts_view
source site coordinates
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...
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 >, 1 > tree_type_1d
nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor< double, PointCloudSearch< view_type > >, PointCloudSearch< view_type >, 2 > tree_type_2d
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...
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 ...
int kdtree_get_point_count() const
Returns the number of source sites.
nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor< double, PointCloudSearch< view_type > >, PointCloudSearch< view_type >, 3 > tree_type_3d
PointCloudSearch(view_type src_pts_view, const local_index_type dimension=-1, const local_index_type max_leaf=-1)
std::shared_ptr< tree_type_3d > _tree_3d
std::shared_ptr< tree_type_2d > _tree_2d
Custom RadiusResultSet for nanoflann that uses pre-allocated space for indices and radii instead of u...
RadiusResultSet(DistanceType radius_, DistanceType *r_dist_, IndexType *i_dist_, const IndexType max_size_)
std::pair< IndexType, DistanceType > worst_item() const
bool addPoint(DistanceType dist, IndexType index)
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...
NeighborLists< view_type > CreateNeighborLists(view_type number_of_neighbors_list)
CreateNeighborLists allows for the construction of an object of type NeighborLists with template dedu...