3 #include <petsc/private/petscimpl.h>
11 …eaves;                // Bit-wise boolean for whether the greater_handle and less_equal_handle ind…
36   PetscKDTreeDestroy - destroy a `PetscKDTree`
41 . tree - tree to destroy
51   PetscCall(PetscFree((*tree)->stems));  in PetscKDTreeDestroy()
52   PetscCall(PetscFree((*tree)->leaves));  in PetscKDTreeDestroy()
53   PetscCall(PetscFree((*tree)->bucket_indices));  in PetscKDTreeDestroy()
54   PetscCall(PetscFree((*tree)->coords_owned));  in PetscKDTreeDestroy()
74   v--;  in RoundToNextPowerOfTwo()
90 // See "Building a Balanced k-d Tree in O(kn log n) Time" https://jcgt.org/published/0004/01/03/
93   const PetscReal *coords = tree->coords;  in PetscKDTreeSortFunc()
94   const PetscInt   dim    = tree->dim;  in PetscKDTreeSortFunc()
97     PetscReal diff = coords[left * dim + axis] - coords[right * dim + axis];  in PetscKDTreeSortFunc()
109 …return PetscKDTreeSortFunc(*(PetscCount *)l, *(PetscCount *)r, kd_ctx->tree, kd_ctx->initial_axis);  in PetscKDTreeTimSort()
114   PetscCount num_coords = tree->num_coords, range_size = end - start, location;  in PetscKDTreeVerifySortedIndices()
122   for (PetscInt d = 1; d < tree->dim; d++) {  in PetscKDTreeVerifySortedIndices()
125 …PetscCheck(location > -1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Sorted indices are not consistent. Cou…  in PetscKDTreeVerifySortedIndices()
140   PetscKDTree tree = kd_build->tree;  in PetscKDTreeBuildStemAndLeaves()
141   PetscInt    dim  = tree->dim;  in PetscKDTreeBuildStemAndLeaves()
145 …if (kd_build->debug_build) PetscCall(PetscKDTreeVerifySortedIndices(tree, sorted_indices, temp, st…  in PetscKDTreeBuildStemAndLeaves()
146   if (end - start <= tree->max_bucket_size) {  in PetscKDTreeBuildStemAndLeaves()
150     PetscCall(PetscSegBufferGetSize(kd_build->leaves, node_handle));  in PetscKDTreeBuildStemAndLeaves()
151     PetscCall(PetscSegBufferGet(kd_build->leaves, 1, &leaf));  in PetscKDTreeBuildStemAndLeaves()
155     PetscCall(PetscIntCast(end - start, &leaf->count));  in PetscKDTreeBuildStemAndLeaves()
156     PetscCall(PetscSegBufferGetSize(kd_build->bucket_indices, &leaf->indices_handle));  in PetscKDTreeBuildStemAndLeaves()
157     PetscCall(PetscSegBufferGet(kd_build->bucket_indices, leaf->count, &bucket_indices));  in PetscKDTreeBuildStemAndLeaves()
158     PetscCall(PetscArraycpy(bucket_indices, &sorted_indices[start], leaf->count));  in PetscKDTreeBuildStemAndLeaves()
159     if (kd_build->copy_coords) {  in PetscKDTreeBuildStemAndLeaves()
161       PetscCall(PetscSegBufferGetSize(kd_build->bucket_coords, &leaf->coords_handle));  in PetscKDTreeBuildStemAndLeaves()
162       PetscCall(PetscSegBufferGet(kd_build->bucket_coords, leaf->count * dim, &bucket_coords));  in PetscKDTreeBuildStemAndLeaves()
163       // Coords are saved in axis-major ordering for better vectorization  in PetscKDTreeBuildStemAndLeaves()
164       for (PetscCount i = 0; i < leaf->count; i++) {  in PetscKDTreeBuildStemAndLeaves()
165 …for (PetscInt d = 0; d < dim; d++) bucket_coords[d * leaf->count + i] = tree->coords[bucket_indice…  in PetscKDTreeBuildStemAndLeaves()
167     } else leaf->coords_handle = -1;  in PetscKDTreeBuildStemAndLeaves()
170     PetscCount num_coords = tree->num_coords;  in PetscKDTreeBuildStemAndLeaves()
173     PetscCount median     = start + PetscCeilInt64(end - start, 2) - 1, lower;  in PetscKDTreeBuildStemAndLeaves()
176     PetscCall(PetscSegBufferGetSize(kd_build->stems, node_handle));  in PetscKDTreeBuildStemAndLeaves()
177     PetscCall(PetscSegBufferGet(kd_build->stems, 1, &stem));  in PetscKDTreeBuildStemAndLeaves()
181     stem->axis = axis;  in PetscKDTreeBuildStemAndLeaves()
183 …stem->split = (tree->coords[tree->dim * median_idx + axis] + tree->coords[tree->dim * medianp1_idx…  in PetscKDTreeBuildStemAndLeaves()
184     PetscCall(PetscArraycpy(temp, &sorted_indices[0 * num_coords + start], end - start));  in PetscKDTreeBuildStemAndLeaves()
186     for (PetscInt d = 1; d < tree->dim; d++) {  in PetscKDTreeBuildStemAndLeaves()
188       lower            = start - 1;  in PetscKDTreeBuildStemAndLeaves()
192           sorted_indices[(d - 1) * num_coords + (++lower)] = sorted_indices[d * num_coords + i];  in PetscKDTreeBuildStemAndLeaves()
194           sorted_indices[(d - 1) * num_coords + (++upper)] = sorted_indices[d * num_coords + i];  in PetscKDTreeBuildStemAndLeaves()
198 …- 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Partitioning into greater bin failed. Range upper bound sho…  in PetscKDTreeBuildStemAndLeaves()
200 …PetscCall(PetscArraycpy(&sorted_indices[(tree->dim - 1) * num_coords + start], temp, end - start));  in PetscKDTreeBuildStemAndLeaves()
202 …sorted_indices, temp, start, lower + 1, depth + 1, &is_less_equal_leaf, &stem->less_equal_handle));  in PetscKDTreeBuildStemAndLeaves()
203     if (is_less_equal_leaf) PetscCall(PetscBTSet(&stem->are_handles_leaves, LESS_EQUAL_BIT));  in PetscKDTreeBuildStemAndLeaves()
204 …_build, sorted_indices, temp, lower + 1, end, depth + 1, &is_greater_leaf, &stem->greater_handle));  in PetscKDTreeBuildStemAndLeaves()
205     if (is_greater_leaf) PetscCall(PetscBTSet(&stem->are_handles_leaves, GREATER_BIT));  in PetscKDTreeBuildStemAndLeaves()
211   PetscKDTreeCreate - create a `PetscKDTree`
216 + num_coords      - number of coordinate points to build the `PetscKDTree`
217 . dim             - the dimension of the coordinates
218 . coords          - array of the coordinates, in point-major order
219 . copy_mode       - behavior handling `coords`, `PETSC_COPY_VALUES` generally more performant
220 - max_bucket_size - maximum number of points stored at each leaf
223 . new_tree - the resulting `PetscKDTree`
228 …PY_VALUES`, the coordinates are copied and organized to optimize vectorization and cache-coherency.
232   Building algorithm detailed in 'Building a Balanced k-d Tree in O(kn log n) Time' Brown, 2015
245 …PetscCheck(num_coords > -1, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "Number of coordinates may not …  in PetscKDTreeCreate()
253   tree->dim             = dim;  in PetscKDTreeCreate()
254   tree->max_bucket_size = max_bucket_size == PETSC_DECIDE ? 32 : max_bucket_size;  in PetscKDTreeCreate()
255   tree->num_coords      = num_coords;  in PetscKDTreeCreate()
259     tree->coords_owned = coords; // fallthrough  in PetscKDTreeCreate()
261     tree->coords = coords;  in PetscKDTreeCreate()
264     PetscCall(PetscMalloc1(num_coords * dim, &tree->coords_owned));  in PetscKDTreeCreate()
265     PetscCall(PetscArraycpy((PetscReal *)tree->coords_owned, coords, num_coords * dim));  in PetscKDTreeCreate()
266     tree->coords = tree->coords_owned;  in PetscKDTreeCreate()
273   kd_ctx->tree = tree;  in PetscKDTreeCreate()
276     kd_ctx->initial_axis = (uint8_t)j;  in PetscKDTreeCreate()
281   PetscInt    num_leaves = (PetscInt)PetscCeilInt64(num_coords, tree->max_bucket_size);  in PetscKDTreeCreate()
285   kd_build->tree        = tree;  in PetscKDTreeCreate()
286   kd_build->copy_coords = copy_mode == PETSC_COPY_VALUES ? PETSC_TRUE : PETSC_FALSE;  in PetscKDTreeCreate()
287   PetscCall(PetscOptionsGetBool(NULL, NULL, "-kdtree_debug", &kd_build->debug_build, NULL));  in PetscKDTreeCreate()
288   PetscCall(PetscSegBufferCreate(sizeof(KDStem), num_stems, &kd_build->stems));  in PetscKDTreeCreate()
289   PetscCall(PetscSegBufferCreate(sizeof(KDLeaf), num_leaves, &kd_build->leaves));  in PetscKDTreeCreate()
290   PetscCall(PetscSegBufferCreate(sizeof(PetscCount), num_coords, &kd_build->bucket_indices));  in PetscKDTreeCreate()
291 …if (kd_build->copy_coords) PetscCall(PetscSegBufferCreate(sizeof(PetscReal), num_coords * dim, &kd…  in PetscKDTreeCreate()
293 …AndLeaves(kd_build, sorted_indices, temp, 0, num_coords, 0, &tree->is_root_leaf, &tree->root_handl…  in PetscKDTreeCreate()
295   PetscCall(PetscSegBufferGetSize(kd_build->stems, &tree->num_stems));  in PetscKDTreeCreate()
296   PetscCall(PetscSegBufferGetSize(kd_build->leaves, &tree->num_leaves));  in PetscKDTreeCreate()
297   PetscCall(PetscSegBufferGetSize(kd_build->bucket_indices, &tree->num_bucket_indices));  in PetscKDTreeCreate()
298   PetscCall(PetscSegBufferExtractAlloc(kd_build->stems, &tree->stems));  in PetscKDTreeCreate()
299   PetscCall(PetscSegBufferExtractAlloc(kd_build->leaves, &tree->leaves));  in PetscKDTreeCreate()
300   PetscCall(PetscSegBufferExtractAlloc(kd_build->bucket_indices, &tree->bucket_indices));  in PetscKDTreeCreate()
301   if (kd_build->copy_coords) {  in PetscKDTreeCreate()
302     PetscCall(PetscFree(tree->coords_owned));  in PetscKDTreeCreate()
303     PetscCall(PetscSegBufferExtractAlloc(kd_build->bucket_coords, &tree->coords_owned));  in PetscKDTreeCreate()
304     tree->coords = tree->coords_owned;  in PetscKDTreeCreate()
305     PetscCall(PetscSegBufferDestroy(&kd_build->bucket_coords));  in PetscKDTreeCreate()
307   PetscCall(PetscSegBufferDestroy(&kd_build->stems));  in PetscKDTreeCreate()
308   PetscCall(PetscSegBufferDestroy(&kd_build->leaves));  in PetscKDTreeCreate()
309   PetscCall(PetscSegBufferDestroy(&kd_build->bucket_indices));  in PetscKDTreeCreate()
319   PetscReal dist = 0;  in PetscSquareDistance()  local
320   for (PetscInt j = 0; j < dim; j++) dist += PetscSqr(x[j] - y[j]);  in PetscSquareDistance()
321   return dist;  in PetscSquareDistance()
326   PetscInt dim = tree->dim;  in PetscKDTreeQueryLeaf()
330   *index        = -1;  in PetscKDTreeQueryLeaf()
332     PetscCount point_index = tree->bucket_indices[leaf.indices_handle + i];  in PetscKDTreeQueryLeaf()
333     PetscReal  dist        = PetscSquareDistance(dim, point, &tree->coords[point_index * dim]);  in PetscKDTreeQueryLeaf()  local
334     if (dist < *distance_sqr) {  in PetscKDTreeQueryLeaf()
335       *distance_sqr = dist;  in PetscKDTreeQueryLeaf()
344   PetscInt dim = tree->dim;  in PetscKDTreeQueryLeaf_CopyCoords()
348   *index        = -1;  in PetscKDTreeQueryLeaf_CopyCoords()
350     // Coord data saved in axis-major ordering for vectorization  in PetscKDTreeQueryLeaf_CopyCoords()
351     PetscReal dist = 0.;  in PetscKDTreeQueryLeaf_CopyCoords()  local
352 …for (PetscInt d = 0; d < dim; d++) dist += PetscSqr(point[d] - tree->coords[leaf.coords_handle + d…  in PetscKDTreeQueryLeaf_CopyCoords()
353     if (dist < *distance_sqr) {  in PetscKDTreeQueryLeaf_CopyCoords()
354       *distance_sqr = dist;  in PetscKDTreeQueryLeaf_CopyCoords()
355       *index        = tree->bucket_indices[leaf.indices_handle + i];  in PetscKDTreeQueryLeaf_CopyCoords()
368     KDLeaf     leaf = tree->leaves[node_handle];  in PetscKDTreeQuery_Recurse()
369     PetscReal  dist;  in PetscKDTreeQuery_Recurse()  local
372 …if (leaf.coords_handle > -1) PetscCall(PetscKDTreeQueryLeaf_CopyCoords(tree, leaf, point, &point_i…  in PetscKDTreeQuery_Recurse()
373     else PetscCall(PetscKDTreeQueryLeaf(tree, leaf, point, &point_index, &dist));  in PetscKDTreeQuery_Recurse()
374     if (dist < *dist_sqr) {  in PetscKDTreeQuery_Recurse()
375       *dist_sqr = dist;  in PetscKDTreeQuery_Recurse()
381   KDStem    stem       = tree->stems[node_handle];  in PetscKDTreeQuery_Recurse()
382   PetscReal old_offset = offset[stem.axis], new_offset = point[stem.axis] - stem.split;  in PetscKDTreeQuery_Recurse()
385     rd += -PetscSqr(old_offset) + PetscSqr(new_offset);  in PetscKDTreeQuery_Recurse()
393     rd += -PetscSqr(old_offset) + PetscSqr(new_offset);  in PetscKDTreeQuery_Recurse()
404   PetscKDTreeQueryPointsNearestNeighbor - find the nearest neighbor in a `PetscKDTree`
409 + tree       - tree to query
410 . num_points - number of points to query
411 . points     - array of the coordinates, in point-major order
412 - tolerance  - tolerance for nearest neighbor
415 + indices   - indices of the nearest neighbor to the query point
416 - distances - distance between the queried point and the nearest neighbor
440   PetscCall(PetscCalloc1(tree->dim, &offsets));  in PetscKDTreeQueryPointsNearestNeighbor()
445     indices[p]   = -1;  in PetscKDTreeQueryPointsNearestNeighbor()
446 …PetscCall(PetscKDTreeQuery_Recurse(tree, &points[p * tree->dim], tree->root_handle, (char)tree->is…  in PetscKDTreeQueryPointsNearestNeighbor()
455   PetscKDTreeView - view a `PetscKDTree`
460 + tree   - tree to view
461 - viewer - visualization context
478   for (PetscCount i = 0; i < tree->num_stems; i++) {  in PetscKDTreeView()
479     KDStem stem = tree->stems[i];  in PetscKDTreeView()
487   for (PetscCount i = 0; i < tree->num_leaves; i++) {  in PetscKDTreeView()
488     KDLeaf leaf = tree->leaves[i];  in PetscKDTreeView()
493       PetscCount bucket_index = tree->bucket_indices[leaf.indices_handle + j];  in PetscKDTreeView()
499       if (leaf.coords_handle > -1) {  in PetscKDTreeView()
500 …for (PetscInt k = 0; k < tree->dim; k++) PetscCall(PetscViewerASCIIPrintf(viewer, "%g ", (double)t…  in PetscKDTreeView()
503 … (PetscInt k = 0; k < tree->dim; k++) PetscCall(PetscViewerASCIIPrintf(viewer, "%g ", (double)tree…  in PetscKDTreeView()