Actual source code: kdtree.c
1: #include <petscis.h>
2: #include <petsc/private/petscimpl.h>
4: // For accessing bitwise Boolean values in are_handles_leaves
5: #define GREATER_BIT 0
6: #define LESS_EQUAL_BIT 1
8: typedef struct {
9: uint8_t axis; // Coordinate direction that stem splits on
10: PetscByte are_handles_leaves; // Bit-wise boolean for whether the greater_handle and less_equal_handle index kdtree->stems or kstree->leaves
11: PetscReal split; // Coordinate value that stem splits on
12: PetscCount greater_handle, less_equal_handle; // Handle (index) of the child nodes of the tree
13: } KDStem;
15: typedef struct {
16: PetscInt count; // Number of points in leaf
17: PetscCount indices_handle, coords_handle; // Index into the indices or coordinates array for the points in leaf
18: } KDLeaf;
20: struct _n_PetscKDTree {
21: PetscInt dim; // Coordinate dimension of the tree
22: PetscInt max_bucket_size; // Maximum number of points stored at each leaf
24: PetscBool is_root_leaf;
25: PetscCount root_handle;
27: PetscCount num_coords, num_leaves, num_stems, num_bucket_indices;
28: const PetscReal *coords, *coords_owned; // Only free owned on Destroy
29: KDLeaf *leaves;
30: KDStem *stems;
31: PetscCount *bucket_indices;
32: };
34: /*@
35: PetscKDTreeDestroy - destroy a `PetscKDTree`
37: Not Collective, No Fortran Support
39: Input Parameters:
40: . tree - tree to destroy
42: Level: advanced
44: .seealso: `PetscKDTree`, `PetscKDTreeCreate()`
45: @*/
46: PetscErrorCode PetscKDTreeDestroy(PetscKDTree *tree)
47: {
48: PetscFunctionBeginUser;
49: if (*tree == NULL) PetscFunctionReturn(PETSC_SUCCESS);
50: PetscCall(PetscFree((*tree)->stems));
51: PetscCall(PetscFree((*tree)->leaves));
52: PetscCall(PetscFree((*tree)->bucket_indices));
53: PetscCall(PetscFree((*tree)->coords_owned));
54: PetscCall(PetscFree(*tree));
55: PetscFunctionReturn(PETSC_SUCCESS);
56: }
58: PetscLogEvent PetscKDTree_Build, PetscKDTree_Query;
59: static PetscErrorCode PetscKDTreeRegisterLogEvents(void)
60: {
61: static PetscBool is_initialized = PETSC_FALSE;
63: PetscFunctionBeginUser;
64: if (is_initialized) PetscFunctionReturn(PETSC_SUCCESS);
65: PetscCall(PetscLogEventRegister("KDTreeBuild", IS_CLASSID, &PetscKDTree_Build));
66: PetscCall(PetscLogEventRegister("KDTreeQuery", IS_CLASSID, &PetscKDTree_Query));
67: PetscFunctionReturn(PETSC_SUCCESS);
68: }
70: // From http://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
71: static inline uint32_t RoundToNextPowerOfTwo(uint32_t v)
72: {
73: v--;
74: v |= v >> 1;
75: v |= v >> 2;
76: v |= v >> 4;
77: v |= v >> 8;
78: v |= v >> 16;
79: v++;
80: return v;
81: }
83: typedef struct {
84: uint8_t initial_axis;
85: PetscKDTree tree;
86: } *KDTreeSortContext;
88: // Sort nodes based on "superkey"
89: // See "Building a Balanced k-d Tree in O(kn log n) Time" https://jcgt.org/published/0004/01/03/
90: static inline int PetscKDTreeSortFunc(PetscCount left, PetscCount right, PetscKDTree tree, uint8_t axis)
91: {
92: const PetscReal *coords = tree->coords;
93: const PetscInt dim = tree->dim;
95: for (PetscInt i = 0; i < dim; i++) {
96: PetscReal diff = coords[left * dim + axis] - coords[right * dim + axis];
97: if (PetscUnlikely(diff == 0)) {
98: axis = (axis + 1) % dim;
99: continue;
100: } else return PetscSign(diff);
101: }
102: return 0; // All components are the same
103: }
105: static int PetscKDTreeTimSort(const void *l, const void *r, PetscCtx ctx)
106: {
107: KDTreeSortContext kd_ctx = (KDTreeSortContext)ctx;
108: return PetscKDTreeSortFunc(*(PetscCount *)l, *(PetscCount *)r, kd_ctx->tree, kd_ctx->initial_axis);
109: }
111: static PetscErrorCode PetscKDTreeVerifySortedIndices(PetscKDTree tree, PetscCount sorted_indices[], PetscCount temp[], PetscCount start, PetscCount end)
112: {
113: PetscCount num_coords = tree->num_coords, range_size = end - start, location;
114: PetscBool has_duplicates;
116: PetscFunctionBeginUser;
117: PetscCall(PetscArraycpy(temp, &sorted_indices[0 * num_coords + start], range_size));
118: PetscCall(PetscSortCount(range_size, temp));
119: PetscCall(PetscSortedCheckDupsCount(range_size, temp, &has_duplicates));
120: PetscCheck(has_duplicates == PETSC_FALSE, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Sorted indices must have unique entries, but found duplicates");
121: for (PetscInt d = 1; d < tree->dim; d++) {
122: for (PetscCount i = start; i < end; i++) {
123: PetscCall(PetscFindCount(sorted_indices[d * num_coords + i], range_size, temp, &location));
124: PetscCheck(location > -1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Sorted indices are not consistent. Could not find %" PetscCount_FMT " from %" PetscInt_FMT " dimensional index in 0th dimension", sorted_indices[d * num_coords + i], d);
125: }
126: }
127: PetscFunctionReturn(PETSC_SUCCESS);
128: }
130: typedef struct {
131: PetscKDTree tree;
132: PetscSegBuffer stems, leaves, bucket_indices, bucket_coords;
133: PetscBool debug_build, copy_coords;
134: } *KDTreeBuild;
136: // The range is end exclusive, so [start,end).
137: static PetscErrorCode PetscKDTreeBuildStemAndLeaves(KDTreeBuild kd_build, PetscCount sorted_indices[], PetscCount temp[], PetscCount start, PetscCount end, PetscInt depth, PetscBool *is_node_leaf, PetscCount *node_handle)
138: {
139: PetscKDTree tree = kd_build->tree;
140: PetscInt dim = tree->dim;
142: PetscFunctionBeginUser;
143: PetscCheck(start <= end, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Start %" PetscCount_FMT " must be less than or equal to end %" PetscCount_FMT, start, end);
144: if (kd_build->debug_build) PetscCall(PetscKDTreeVerifySortedIndices(tree, sorted_indices, temp, start, end));
145: if (end - start <= tree->max_bucket_size) {
146: KDLeaf *leaf;
147: PetscCount *bucket_indices;
149: PetscCall(PetscSegBufferGetSize(kd_build->leaves, node_handle));
150: PetscCall(PetscSegBufferGet(kd_build->leaves, 1, &leaf));
151: PetscCall(PetscMemzero(leaf, sizeof(KDLeaf)));
152: *is_node_leaf = PETSC_TRUE;
154: PetscCall(PetscIntCast(end - start, &leaf->count));
155: PetscCall(PetscSegBufferGetSize(kd_build->bucket_indices, &leaf->indices_handle));
156: PetscCall(PetscSegBufferGet(kd_build->bucket_indices, leaf->count, &bucket_indices));
157: PetscCall(PetscArraycpy(bucket_indices, &sorted_indices[start], leaf->count));
158: if (kd_build->copy_coords) {
159: PetscReal *bucket_coords;
160: PetscCall(PetscSegBufferGetSize(kd_build->bucket_coords, &leaf->coords_handle));
161: PetscCall(PetscSegBufferGet(kd_build->bucket_coords, leaf->count * dim, &bucket_coords));
162: // Coords are saved in axis-major ordering for better vectorization
163: for (PetscCount i = 0; i < leaf->count; i++) {
164: for (PetscInt d = 0; d < dim; d++) bucket_coords[d * leaf->count + i] = tree->coords[bucket_indices[i] * dim + d];
165: }
166: } else leaf->coords_handle = -1;
167: } else {
168: KDStem *stem;
169: PetscCount num_coords = tree->num_coords;
170: uint8_t axis = (uint8_t)(depth % dim);
171: PetscBool is_greater_leaf, is_less_equal_leaf;
172: PetscCount median = start + PetscCeilInt64(end - start, 2) - 1, lower;
173: PetscCount median_idx = sorted_indices[median], medianp1_idx = sorted_indices[median + 1];
175: PetscCall(PetscSegBufferGetSize(kd_build->stems, node_handle));
176: PetscCall(PetscSegBufferGet(kd_build->stems, 1, &stem));
177: PetscCall(PetscMemzero(stem, sizeof(KDStem)));
178: *is_node_leaf = PETSC_FALSE;
180: stem->axis = axis;
181: // Place split halfway between the "boundary" nodes of the partitioning
182: stem->split = (tree->coords[tree->dim * median_idx + axis] + tree->coords[tree->dim * medianp1_idx + axis]) / 2;
183: PetscCall(PetscArraycpy(temp, &sorted_indices[0 * num_coords + start], end - start));
184: lower = median; // Set lower in case dim == 1
185: for (PetscInt d = 1; d < tree->dim; d++) {
186: PetscCount upper = median;
187: lower = start - 1;
188: for (PetscCount i = start; i < end; i++) {
189: // In case of duplicate coord point equal to the median coord point, limit lower partition to median, ensuring balanced tree
190: if (lower < median && PetscKDTreeSortFunc(sorted_indices[d * num_coords + i], median_idx, tree, axis) <= 0) {
191: sorted_indices[(d - 1) * num_coords + (++lower)] = sorted_indices[d * num_coords + i];
192: } else {
193: sorted_indices[(d - 1) * num_coords + (++upper)] = sorted_indices[d * num_coords + i];
194: }
195: }
196: PetscCheck(lower == median, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Partitioning into less_equal bin failed. Range upper bound should be %" PetscCount_FMT " but partitioning resulted in %" PetscCount_FMT, median, lower);
197: PetscCheck(upper == end - 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Partitioning into greater bin failed. Range upper bound should be %" PetscCount_FMT " but partitioning resulted in %" PetscCount_FMT, upper, end - 1);
198: }
199: PetscCall(PetscArraycpy(&sorted_indices[(tree->dim - 1) * num_coords + start], temp, end - start));
201: PetscCall(PetscKDTreeBuildStemAndLeaves(kd_build, sorted_indices, temp, start, lower + 1, depth + 1, &is_less_equal_leaf, &stem->less_equal_handle));
202: if (is_less_equal_leaf) PetscCall(PetscBTSet(&stem->are_handles_leaves, LESS_EQUAL_BIT));
203: PetscCall(PetscKDTreeBuildStemAndLeaves(kd_build, sorted_indices, temp, lower + 1, end, depth + 1, &is_greater_leaf, &stem->greater_handle));
204: if (is_greater_leaf) PetscCall(PetscBTSet(&stem->are_handles_leaves, GREATER_BIT));
205: }
206: PetscFunctionReturn(PETSC_SUCCESS);
207: }
209: /*@
210: PetscKDTreeCreate - create a `PetscKDTree`
212: Not Collective, No Fortran Support
214: Input Parameters:
215: + num_coords - number of coordinate points to build the `PetscKDTree`
216: . dim - the dimension of the coordinates
217: . coords - array of the coordinates, in point-major order
218: . copy_mode - behavior handling `coords`, `PETSC_COPY_VALUES` generally more performant
219: - max_bucket_size - maximum number of points stored at each leaf
221: Output Parameter:
222: . new_tree - the resulting `PetscKDTree`
224: Level: advanced
226: Note:
227: When `copy_mode == PETSC_COPY_VALUES`, the coordinates are copied and organized to optimize vectorization and cache-coherency.
228: It is recommended to run this way if the extra memory use is not a concern and it has very little impact on the `PetscKDTree` creation time.
230: Developer Note:
231: Building algorithm detailed in 'Building a Balanced k-d Tree in O(kn log n) Time' Brown, 2015
233: .seealso: `PetscKDTree`, `PetscKDTreeDestroy()`, `PetscKDTreeQueryPointsNearestNeighbor()`
234: @*/
235: PetscErrorCode PetscKDTreeCreate(PetscCount num_coords, PetscInt dim, const PetscReal coords[], PetscCopyMode copy_mode, PetscInt max_bucket_size, PetscKDTree *new_tree)
236: {
237: PetscKDTree tree;
238: PetscCount *sorted_indices, *temp;
240: PetscFunctionBeginUser;
241: PetscCall(PetscKDTreeRegisterLogEvents());
242: PetscCall(PetscLogEventBegin(PetscKDTree_Build, 0, 0, 0, 0));
243: PetscCheck(dim > 0, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "Dimension of PetscKDTree must be greater than 0, received %" PetscInt_FMT, dim);
244: PetscCheck(num_coords > -1, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "Number of coordinates may not be negative, received %" PetscCount_FMT, num_coords);
245: if (num_coords == 0) {
246: *new_tree = NULL;
247: PetscFunctionReturn(PETSC_SUCCESS);
248: }
249: PetscAssertPointer(coords, 3);
250: PetscAssertPointer(new_tree, 6);
251: PetscCall(PetscNew(&tree));
252: tree->dim = dim;
253: tree->max_bucket_size = max_bucket_size == PETSC_DECIDE ? 32 : max_bucket_size;
254: tree->num_coords = num_coords;
256: switch (copy_mode) {
257: case PETSC_OWN_POINTER:
258: tree->coords_owned = coords; // fallthrough
259: case PETSC_USE_POINTER:
260: tree->coords = coords;
261: break;
262: case PETSC_COPY_VALUES:
263: PetscCall(PetscMalloc1(num_coords * dim, &tree->coords_owned));
264: PetscCall(PetscArraycpy((PetscReal *)tree->coords_owned, coords, num_coords * dim));
265: tree->coords = tree->coords_owned;
266: break;
267: }
269: KDTreeSortContext kd_ctx;
270: PetscCall(PetscMalloc2(num_coords * dim, &sorted_indices, num_coords, &temp));
271: PetscCall(PetscNew(&kd_ctx));
272: kd_ctx->tree = tree;
273: for (PetscInt j = 0; j < dim; j++) {
274: for (PetscCount i = 0; i < num_coords; i++) sorted_indices[num_coords * j + i] = i;
275: kd_ctx->initial_axis = (uint8_t)j;
276: PetscCall(PetscTimSort((PetscInt)num_coords, &sorted_indices[num_coords * j], sizeof(*sorted_indices), PetscKDTreeTimSort, kd_ctx));
277: }
278: PetscCall(PetscFree(kd_ctx));
280: PetscInt num_leaves = (PetscInt)PetscCeilInt64(num_coords, tree->max_bucket_size);
281: PetscInt num_stems = RoundToNextPowerOfTwo((uint32_t)num_leaves);
282: KDTreeBuild kd_build;
283: PetscCall(PetscNew(&kd_build));
284: kd_build->tree = tree;
285: kd_build->copy_coords = copy_mode == PETSC_COPY_VALUES ? PETSC_TRUE : PETSC_FALSE;
286: PetscCall(PetscOptionsGetBool(NULL, NULL, "-kdtree_debug", &kd_build->debug_build, NULL));
287: PetscCall(PetscSegBufferCreate(sizeof(KDStem), num_stems, &kd_build->stems));
288: PetscCall(PetscSegBufferCreate(sizeof(KDLeaf), num_leaves, &kd_build->leaves));
289: PetscCall(PetscSegBufferCreate(sizeof(PetscCount), num_coords, &kd_build->bucket_indices));
290: if (kd_build->copy_coords) PetscCall(PetscSegBufferCreate(sizeof(PetscReal), num_coords * dim, &kd_build->bucket_coords));
292: PetscCall(PetscKDTreeBuildStemAndLeaves(kd_build, sorted_indices, temp, 0, num_coords, 0, &tree->is_root_leaf, &tree->root_handle));
294: PetscCall(PetscSegBufferGetSize(kd_build->stems, &tree->num_stems));
295: PetscCall(PetscSegBufferGetSize(kd_build->leaves, &tree->num_leaves));
296: PetscCall(PetscSegBufferGetSize(kd_build->bucket_indices, &tree->num_bucket_indices));
297: PetscCall(PetscSegBufferExtractAlloc(kd_build->stems, &tree->stems));
298: PetscCall(PetscSegBufferExtractAlloc(kd_build->leaves, &tree->leaves));
299: PetscCall(PetscSegBufferExtractAlloc(kd_build->bucket_indices, &tree->bucket_indices));
300: if (kd_build->copy_coords) {
301: PetscCall(PetscFree(tree->coords_owned));
302: PetscCall(PetscSegBufferExtractAlloc(kd_build->bucket_coords, &tree->coords_owned));
303: tree->coords = tree->coords_owned;
304: PetscCall(PetscSegBufferDestroy(&kd_build->bucket_coords));
305: }
306: PetscCall(PetscSegBufferDestroy(&kd_build->stems));
307: PetscCall(PetscSegBufferDestroy(&kd_build->leaves));
308: PetscCall(PetscSegBufferDestroy(&kd_build->bucket_indices));
309: PetscCall(PetscFree(kd_build));
310: PetscCall(PetscFree2(sorted_indices, temp));
311: *new_tree = tree;
312: PetscCall(PetscLogEventEnd(PetscKDTree_Build, 0, 0, 0, 0));
313: PetscFunctionReturn(PETSC_SUCCESS);
314: }
316: static inline PetscReal PetscSquareDistance(PetscInt dim, const PetscReal *PETSC_RESTRICT x, const PetscReal *PETSC_RESTRICT y)
317: {
318: PetscReal dist = 0;
319: for (PetscInt j = 0; j < dim; j++) dist += PetscSqr(x[j] - y[j]);
320: return dist;
321: }
323: static inline PetscErrorCode PetscKDTreeQueryLeaf(PetscKDTree tree, KDLeaf leaf, const PetscReal point[], PetscCount *index, PetscReal *distance_sqr)
324: {
325: PetscInt dim = tree->dim;
327: PetscFunctionBeginUser;
328: *distance_sqr = PETSC_MAX_REAL;
329: *index = -1;
330: for (PetscInt i = 0; i < leaf.count; i++) {
331: PetscCount point_index = tree->bucket_indices[leaf.indices_handle + i];
332: PetscReal dist = PetscSquareDistance(dim, point, &tree->coords[point_index * dim]);
333: if (dist < *distance_sqr) {
334: *distance_sqr = dist;
335: *index = point_index;
336: }
337: }
338: PetscFunctionReturn(PETSC_SUCCESS);
339: }
341: static inline PetscErrorCode PetscKDTreeQueryLeaf_CopyCoords(PetscKDTree tree, KDLeaf leaf, const PetscReal point[], PetscCount *index, PetscReal *distance_sqr)
342: {
343: PetscInt dim = tree->dim;
345: PetscFunctionBeginUser;
346: *distance_sqr = PETSC_MAX_REAL;
347: *index = -1;
348: for (PetscInt i = 0; i < leaf.count; i++) {
349: // Coord data saved in axis-major ordering for vectorization
350: PetscReal dist = 0.;
351: for (PetscInt d = 0; d < dim; d++) dist += PetscSqr(point[d] - tree->coords[leaf.coords_handle + d * leaf.count + i]);
352: if (dist < *distance_sqr) {
353: *distance_sqr = dist;
354: *index = tree->bucket_indices[leaf.indices_handle + i];
355: }
356: }
357: PetscFunctionReturn(PETSC_SUCCESS);
358: }
360: // Recursive point query from 'Algorithms for Fast Vector Quantization' by Sunil Arya and David Mount
361: // Variant also implemented in pykdtree
362: static PetscErrorCode PetscKDTreeQuery_Recurse(PetscKDTree tree, const PetscReal point[], PetscCount node_handle, char is_node_leaf, PetscReal offset[], PetscReal rd, PetscReal tol_sqr, PetscCount *index, PetscReal *dist_sqr)
363: {
364: PetscFunctionBeginUser;
365: if (*dist_sqr < tol_sqr) PetscFunctionReturn(PETSC_SUCCESS);
366: if (is_node_leaf) {
367: KDLeaf leaf = tree->leaves[node_handle];
368: PetscReal dist;
369: PetscCount point_index;
371: if (leaf.coords_handle > -1) PetscCall(PetscKDTreeQueryLeaf_CopyCoords(tree, leaf, point, &point_index, &dist));
372: else PetscCall(PetscKDTreeQueryLeaf(tree, leaf, point, &point_index, &dist));
373: if (dist < *dist_sqr) {
374: *dist_sqr = dist;
375: *index = point_index;
376: }
377: PetscFunctionReturn(PETSC_SUCCESS);
378: }
380: KDStem stem = tree->stems[node_handle];
381: PetscReal old_offset = offset[stem.axis], new_offset = point[stem.axis] - stem.split;
382: if (new_offset <= 0) {
383: PetscCall(PetscKDTreeQuery_Recurse(tree, point, stem.less_equal_handle, PetscBTLookup(&stem.are_handles_leaves, LESS_EQUAL_BIT), offset, rd, tol_sqr, index, dist_sqr));
384: rd += -PetscSqr(old_offset) + PetscSqr(new_offset);
385: if (rd < *dist_sqr) {
386: offset[stem.axis] = new_offset;
387: PetscCall(PetscKDTreeQuery_Recurse(tree, point, stem.greater_handle, PetscBTLookup(&stem.are_handles_leaves, GREATER_BIT), offset, rd, tol_sqr, index, dist_sqr));
388: offset[stem.axis] = old_offset;
389: }
390: } else {
391: PetscCall(PetscKDTreeQuery_Recurse(tree, point, stem.greater_handle, PetscBTLookup(&stem.are_handles_leaves, GREATER_BIT), offset, rd, tol_sqr, index, dist_sqr));
392: rd += -PetscSqr(old_offset) + PetscSqr(new_offset);
393: if (rd < *dist_sqr) {
394: offset[stem.axis] = new_offset;
395: PetscCall(PetscKDTreeQuery_Recurse(tree, point, stem.less_equal_handle, PetscBTLookup(&stem.are_handles_leaves, LESS_EQUAL_BIT), offset, rd, tol_sqr, index, dist_sqr));
396: offset[stem.axis] = old_offset;
397: }
398: }
399: PetscFunctionReturn(PETSC_SUCCESS);
400: }
402: /*@
403: PetscKDTreeQueryPointsNearestNeighbor - find the nearest neighbor in a `PetscKDTree`
405: Not Collective, No Fortran Support
407: Input Parameters:
408: + tree - tree to query
409: . num_points - number of points to query
410: . points - array of the coordinates, in point-major order
411: - tolerance - tolerance for nearest neighbor
413: Output Parameters:
414: + indices - indices of the nearest neighbor to the query point
415: - distances - distance between the queried point and the nearest neighbor
417: Level: advanced
419: Notes:
420: When traversing the tree, if a point has been found to be closer than the `tolerance`, the function short circuits and doesn't check for any closer points.
422: The `indices` and `distances` arrays should be at least of size `num_points`.
424: .seealso: `PetscKDTree`, `PetscKDTreeCreate()`
425: @*/
426: PetscErrorCode PetscKDTreeQueryPointsNearestNeighbor(PetscKDTree tree, PetscCount num_points, const PetscReal points[], PetscReal tolerance, PetscCount indices[], PetscReal distances[])
427: {
428: PetscReal *offsets, rd;
430: PetscFunctionBeginUser;
431: PetscCall(PetscLogEventBegin(PetscKDTree_Query, 0, 0, 0, 0));
432: if (tree == NULL) {
433: PetscCheck(num_points == 0, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "num_points may only be zero, if tree is NULL");
434: PetscFunctionReturn(PETSC_SUCCESS);
435: }
436: PetscAssertPointer(points, 3);
437: PetscAssertPointer(indices, 5);
438: PetscAssertPointer(distances, 6);
439: PetscCall(PetscCalloc1(tree->dim, &offsets));
441: for (PetscCount p = 0; p < num_points; p++) {
442: rd = 0.;
443: distances[p] = PETSC_MAX_REAL;
444: indices[p] = -1;
445: PetscCall(PetscKDTreeQuery_Recurse(tree, &points[p * tree->dim], tree->root_handle, (char)tree->is_root_leaf, offsets, rd, PetscSqr(tolerance), &indices[p], &distances[p]));
446: distances[p] = PetscSqrtReal(distances[p]);
447: }
448: PetscCall(PetscFree(offsets));
449: PetscCall(PetscLogEventEnd(PetscKDTree_Query, 0, 0, 0, 0));
450: PetscFunctionReturn(PETSC_SUCCESS);
451: }
453: /*@
454: PetscKDTreeView - view a `PetscKDTree`
456: Not Collective, No Fortran Support
458: Input Parameters:
459: + tree - tree to view
460: - viewer - visualization context
462: Level: advanced
464: .seealso: `PetscKDTree`, `PetscKDTreeCreate()`, `PetscViewer`
465: @*/
466: PetscErrorCode PetscKDTreeView(PetscKDTree tree, PetscViewer viewer)
467: {
468: PetscFunctionBeginUser;
470: else PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_SELF, &viewer));
471: if (tree == NULL) PetscFunctionReturn(PETSC_SUCCESS);
473: PetscCall(PetscViewerASCIIPrintf(viewer, "KDTree:\n"));
474: PetscCall(PetscViewerASCIIPushTab(viewer)); // KDTree:
475: PetscCall(PetscViewerASCIIPrintf(viewer, "Stems:\n"));
476: PetscCall(PetscViewerASCIIPushTab(viewer)); // Stems:
477: for (PetscCount i = 0; i < tree->num_stems; i++) {
478: KDStem stem = tree->stems[i];
479: PetscCall(PetscViewerASCIIPrintf(viewer, "Stem %" PetscCount_FMT ": Axis=%" PetscInt_FMT " Split=%g Greater_%s=%" PetscCount_FMT " Lesser_Equal_%s=%" PetscCount_FMT "\n", i, (PetscInt)stem.axis, (double)stem.split, PetscBTLookup(&stem.are_handles_leaves, GREATER_BIT) ? "Leaf" : "Stem",
480: stem.greater_handle, PetscBTLookup(&stem.are_handles_leaves, LESS_EQUAL_BIT) ? "Leaf" : "Stem", stem.less_equal_handle));
481: }
482: PetscCall(PetscViewerASCIIPopTab(viewer)); // Stems:
484: PetscCall(PetscViewerASCIIPrintf(viewer, "Leaves:\n"));
485: PetscCall(PetscViewerASCIIPushTab(viewer)); // Leaves:
486: for (PetscCount i = 0; i < tree->num_leaves; i++) {
487: KDLeaf leaf = tree->leaves[i];
488: PetscCall(PetscViewerASCIIPrintf(viewer, "Leaf %" PetscCount_FMT ": Count=%" PetscInt_FMT, i, leaf.count));
489: PetscCall(PetscViewerASCIIPushTab(viewer)); // Coords:
490: for (PetscInt j = 0; j < leaf.count; j++) {
491: PetscInt tabs;
492: PetscCount bucket_index = tree->bucket_indices[leaf.indices_handle + j];
493: PetscCall(PetscViewerASCIIPrintf(viewer, "\n"));
494: PetscCall(PetscViewerASCIIPrintf(viewer, "%" PetscCount_FMT ": ", bucket_index));
496: PetscCall(PetscViewerASCIIGetTab(viewer, &tabs));
497: PetscCall(PetscViewerASCIISetTab(viewer, 0));
498: if (leaf.coords_handle > -1) {
499: for (PetscInt k = 0; k < tree->dim; k++) PetscCall(PetscViewerASCIIPrintf(viewer, "%g ", (double)tree->coords[leaf.coords_handle + leaf.count * k + j]));
500: PetscCall(PetscViewerASCIIPrintf(viewer, " (stored at leaf)"));
501: } else {
502: for (PetscInt k = 0; k < tree->dim; k++) PetscCall(PetscViewerASCIIPrintf(viewer, "%g ", (double)tree->coords[bucket_index * tree->dim + k]));
503: }
504: PetscCall(PetscViewerASCIISetTab(viewer, tabs));
505: }
506: PetscCall(PetscViewerASCIIPopTab(viewer)); // Coords:
507: PetscCall(PetscViewerASCIIPrintf(viewer, "\n"));
508: }
509: PetscCall(PetscViewerASCIIPopTab(viewer)); // Leaves:
510: PetscCall(PetscViewerASCIIPopTab(viewer)); // KDTree:
511: PetscFunctionReturn(PETSC_SUCCESS);
512: }