Actual source code: kdtree.c

      1: #include <petscviewer.h>
      2: #include <petscis.h>
      3: #include <petsc/private/petscimpl.h>

      5: // For accessing bitwise Boolean values in are_handles_leaves
      6: #define GREATER_BIT    0
      7: #define LESS_EQUAL_BIT 1

      9: typedef struct {
     10:   uint8_t    axis;                              // Coordinate direction that stem splits on
     11:   PetscByte  are_handles_leaves;                // Bit-wise boolean for whether the greater_handle and less_equal_handle index kdtree->stems or kstree->leaves
     12:   PetscReal  split;                             // Coordinate value that stem splits on
     13:   PetscCount greater_handle, less_equal_handle; // Handle (index) of the child nodes of the tree
     14: } KDStem;

     16: typedef struct {
     17:   PetscInt   count;                         // Number of points in leaf
     18:   PetscCount indices_handle, coords_handle; // Index into the indices or coordinates array for the points in leaf
     19: } KDLeaf;

     21: struct _n_PetscKDTree {
     22:   PetscInt dim;             // Coordinate dimension of the tree
     23:   PetscInt max_bucket_size; // Maximum number of points stored at each leaf

     25:   PetscBool  is_root_leaf;
     26:   PetscCount root_handle;

     28:   PetscCount       num_coords, num_leaves, num_stems, num_bucket_indices;
     29:   const PetscReal *coords, *coords_owned; // Only free owned on Destroy
     30:   KDLeaf          *leaves;
     31:   KDStem          *stems;
     32:   PetscCount      *bucket_indices;
     33: };

     35: /*@C
     36:   PetscKDTreeDestroy - destroy a `PetscKDTree`

     38:   Not Collective, No Fortran Support

     40:   Input Parameters:
     41: . tree - tree to destroy

     43:   Level: advanced

     45: .seealso: `PetscKDTree`, `PetscKDTreeCreate()`
     46: @*/
     47: PetscErrorCode PetscKDTreeDestroy(PetscKDTree *tree)
     48: {
     49:   PetscFunctionBeginUser;
     50:   if (*tree == NULL) PetscFunctionReturn(PETSC_SUCCESS);
     51:   PetscCall(PetscFree((*tree)->stems));
     52:   PetscCall(PetscFree((*tree)->leaves));
     53:   PetscCall(PetscFree((*tree)->bucket_indices));
     54:   PetscCall(PetscFree((*tree)->coords_owned));
     55:   PetscCall(PetscFree(*tree));
     56:   PetscFunctionReturn(PETSC_SUCCESS);
     57: }

     59: PetscLogEvent         PetscKDTree_Build, PetscKDTree_Query;
     60: static PetscErrorCode PetscKDTreeRegisterLogEvents(void)
     61: {
     62:   static PetscBool is_initialized = PETSC_FALSE;

     64:   PetscFunctionBeginUser;
     65:   if (is_initialized) PetscFunctionReturn(PETSC_SUCCESS);
     66:   PetscCall(PetscLogEventRegister("KDTreeBuild", IS_CLASSID, &PetscKDTree_Build));
     67:   PetscCall(PetscLogEventRegister("KDTreeQuery", IS_CLASSID, &PetscKDTree_Query));
     68:   PetscFunctionReturn(PETSC_SUCCESS);
     69: }

     71: // From http://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
     72: static inline uint32_t RoundToNextPowerOfTwo(uint32_t v)
     73: {
     74:   v--;
     75:   v |= v >> 1;
     76:   v |= v >> 2;
     77:   v |= v >> 4;
     78:   v |= v >> 8;
     79:   v |= v >> 16;
     80:   v++;
     81:   return v;
     82: }

     84: typedef struct {
     85:   uint8_t     initial_axis;
     86:   PetscKDTree tree;
     87: } *KDTreeSortContext;

     89: // Sort nodes based on "superkey"
     90: // See "Building a Balanced k-d Tree in O(kn log n) Time" https://jcgt.org/published/0004/01/03/
     91: static inline int PetscKDTreeSortFunc(PetscCount left, PetscCount right, PetscKDTree tree, uint8_t axis)
     92: {
     93:   const PetscReal *coords = tree->coords;
     94:   const PetscInt   dim    = tree->dim;

     96:   for (PetscInt i = 0; i < dim; i++) {
     97:     PetscReal diff = coords[left * dim + axis] - coords[right * dim + axis];
     98:     if (PetscUnlikely(diff == 0)) {
     99:       axis = (axis + 1) % dim;
    100:       continue;
    101:     } else return PetscSign(diff);
    102:   }
    103:   return 0; // All components are the same
    104: }

    106: static int PetscKDTreeTimSort(const void *l, const void *r, PetscCtx ctx)
    107: {
    108:   KDTreeSortContext kd_ctx = (KDTreeSortContext)ctx;
    109:   return PetscKDTreeSortFunc(*(PetscCount *)l, *(PetscCount *)r, kd_ctx->tree, kd_ctx->initial_axis);
    110: }

    112: static PetscErrorCode PetscKDTreeVerifySortedIndices(PetscKDTree tree, PetscCount sorted_indices[], PetscCount temp[], PetscCount start, PetscCount end)
    113: {
    114:   PetscCount num_coords = tree->num_coords, range_size = end - start, location;
    115:   PetscBool  has_duplicates;

    117:   PetscFunctionBeginUser;
    118:   PetscCall(PetscArraycpy(temp, &sorted_indices[0 * num_coords + start], range_size));
    119:   PetscCall(PetscSortCount(range_size, temp));
    120:   PetscCall(PetscSortedCheckDupsCount(range_size, temp, &has_duplicates));
    121:   PetscCheck(has_duplicates == PETSC_FALSE, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Sorted indices must have unique entries, but found duplicates");
    122:   for (PetscInt d = 1; d < tree->dim; d++) {
    123:     for (PetscCount i = start; i < end; i++) {
    124:       PetscCall(PetscFindCount(sorted_indices[d * num_coords + i], range_size, temp, &location));
    125:       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);
    126:     }
    127:   }
    128:   PetscFunctionReturn(PETSC_SUCCESS);
    129: }

    131: typedef struct {
    132:   PetscKDTree    tree;
    133:   PetscSegBuffer stems, leaves, bucket_indices, bucket_coords;
    134:   PetscBool      debug_build, copy_coords;
    135: } *KDTreeBuild;

    137: // The range is end exclusive, so [start,end).
    138: static PetscErrorCode PetscKDTreeBuildStemAndLeaves(KDTreeBuild kd_build, PetscCount sorted_indices[], PetscCount temp[], PetscCount start, PetscCount end, PetscInt depth, PetscBool *is_node_leaf, PetscCount *node_handle)
    139: {
    140:   PetscKDTree tree = kd_build->tree;
    141:   PetscInt    dim  = tree->dim;

    143:   PetscFunctionBeginUser;
    144:   PetscCheck(start <= end, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Start %" PetscCount_FMT " must be less than or equal to end %" PetscCount_FMT, start, end);
    145:   if (kd_build->debug_build) PetscCall(PetscKDTreeVerifySortedIndices(tree, sorted_indices, temp, start, end));
    146:   if (end - start <= tree->max_bucket_size) {
    147:     KDLeaf     *leaf;
    148:     PetscCount *bucket_indices;

    150:     PetscCall(PetscSegBufferGetSize(kd_build->leaves, node_handle));
    151:     PetscCall(PetscSegBufferGet(kd_build->leaves, 1, &leaf));
    152:     PetscCall(PetscMemzero(leaf, sizeof(KDLeaf)));
    153:     *is_node_leaf = PETSC_TRUE;

    155:     PetscCall(PetscIntCast(end - start, &leaf->count));
    156:     PetscCall(PetscSegBufferGetSize(kd_build->bucket_indices, &leaf->indices_handle));
    157:     PetscCall(PetscSegBufferGet(kd_build->bucket_indices, leaf->count, &bucket_indices));
    158:     PetscCall(PetscArraycpy(bucket_indices, &sorted_indices[start], leaf->count));
    159:     if (kd_build->copy_coords) {
    160:       PetscReal *bucket_coords;
    161:       PetscCall(PetscSegBufferGetSize(kd_build->bucket_coords, &leaf->coords_handle));
    162:       PetscCall(PetscSegBufferGet(kd_build->bucket_coords, leaf->count * dim, &bucket_coords));
    163:       // Coords are saved in axis-major ordering for better vectorization
    164:       for (PetscCount i = 0; i < leaf->count; i++) {
    165:         for (PetscInt d = 0; d < dim; d++) bucket_coords[d * leaf->count + i] = tree->coords[bucket_indices[i] * dim + d];
    166:       }
    167:     } else leaf->coords_handle = -1;
    168:   } else {
    169:     KDStem    *stem;
    170:     PetscCount num_coords = tree->num_coords;
    171:     uint8_t    axis       = (uint8_t)(depth % dim);
    172:     PetscBool  is_greater_leaf, is_less_equal_leaf;
    173:     PetscCount median     = start + PetscCeilInt64(end - start, 2) - 1, lower;
    174:     PetscCount median_idx = sorted_indices[median], medianp1_idx = sorted_indices[median + 1];

    176:     PetscCall(PetscSegBufferGetSize(kd_build->stems, node_handle));
    177:     PetscCall(PetscSegBufferGet(kd_build->stems, 1, &stem));
    178:     PetscCall(PetscMemzero(stem, sizeof(KDStem)));
    179:     *is_node_leaf = PETSC_FALSE;

    181:     stem->axis = axis;
    182:     // Place split halfway between the "boundary" nodes of the partitioning
    183:     stem->split = (tree->coords[tree->dim * median_idx + axis] + tree->coords[tree->dim * medianp1_idx + axis]) / 2;
    184:     PetscCall(PetscArraycpy(temp, &sorted_indices[0 * num_coords + start], end - start));
    185:     lower = median; // Set lower in case dim == 1
    186:     for (PetscInt d = 1; d < tree->dim; d++) {
    187:       PetscCount upper = median;
    188:       lower            = start - 1;
    189:       for (PetscCount i = start; i < end; i++) {
    190:         // In case of duplicate coord point equal to the median coord point, limit lower partition to median, ensuring balanced tree
    191:         if (lower < median && PetscKDTreeSortFunc(sorted_indices[d * num_coords + i], median_idx, tree, axis) <= 0) {
    192:           sorted_indices[(d - 1) * num_coords + (++lower)] = sorted_indices[d * num_coords + i];
    193:         } else {
    194:           sorted_indices[(d - 1) * num_coords + (++upper)] = sorted_indices[d * num_coords + i];
    195:         }
    196:       }
    197:       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);
    198:       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);
    199:     }
    200:     PetscCall(PetscArraycpy(&sorted_indices[(tree->dim - 1) * num_coords + start], temp, end - start));

    202:     PetscCall(PetscKDTreeBuildStemAndLeaves(kd_build, sorted_indices, temp, start, lower + 1, depth + 1, &is_less_equal_leaf, &stem->less_equal_handle));
    203:     if (is_less_equal_leaf) PetscCall(PetscBTSet(&stem->are_handles_leaves, LESS_EQUAL_BIT));
    204:     PetscCall(PetscKDTreeBuildStemAndLeaves(kd_build, sorted_indices, temp, lower + 1, end, depth + 1, &is_greater_leaf, &stem->greater_handle));
    205:     if (is_greater_leaf) PetscCall(PetscBTSet(&stem->are_handles_leaves, GREATER_BIT));
    206:   }
    207:   PetscFunctionReturn(PETSC_SUCCESS);
    208: }

    210: /*@C
    211:   PetscKDTreeCreate - create a `PetscKDTree`

    213:   Not Collective, No Fortran Support

    215:   Input Parameters:
    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

    222:   Output Parameter:
    223: . new_tree - the resulting `PetscKDTree`

    225:   Level: advanced

    227:   Note:
    228:   When `copy_mode == PETSC_COPY_VALUES`, the coordinates are copied and organized to optimize vectorization and cache-coherency.
    229:   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.

    231:   Developer Note:
    232:   Building algorithm detailed in 'Building a Balanced k-d Tree in O(kn log n) Time' Brown, 2015

    234: .seealso: `PetscKDTree`, `PetscKDTreeDestroy()`, `PetscKDTreeQueryPointsNearestNeighbor()`
    235: @*/
    236: PetscErrorCode PetscKDTreeCreate(PetscCount num_coords, PetscInt dim, const PetscReal coords[], PetscCopyMode copy_mode, PetscInt max_bucket_size, PetscKDTree *new_tree)
    237: {
    238:   PetscKDTree tree;
    239:   PetscCount *sorted_indices, *temp;

    241:   PetscFunctionBeginUser;
    242:   PetscCall(PetscKDTreeRegisterLogEvents());
    243:   PetscCall(PetscLogEventBegin(PetscKDTree_Build, 0, 0, 0, 0));
    244:   PetscCheck(dim > 0, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "Dimension of PetscKDTree must be greater than 0, received %" PetscInt_FMT, dim);
    245:   PetscCheck(num_coords > -1, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "Number of coordinates may not be negative, received %" PetscCount_FMT, num_coords);
    246:   if (num_coords == 0) {
    247:     *new_tree = NULL;
    248:     PetscFunctionReturn(PETSC_SUCCESS);
    249:   }
    250:   PetscAssertPointer(coords, 3);
    251:   PetscAssertPointer(new_tree, 6);
    252:   PetscCall(PetscNew(&tree));
    253:   tree->dim             = dim;
    254:   tree->max_bucket_size = max_bucket_size == PETSC_DECIDE ? 32 : max_bucket_size;
    255:   tree->num_coords      = num_coords;

    257:   switch (copy_mode) {
    258:   case PETSC_OWN_POINTER:
    259:     tree->coords_owned = coords; // fallthrough
    260:   case PETSC_USE_POINTER:
    261:     tree->coords = coords;
    262:     break;
    263:   case PETSC_COPY_VALUES:
    264:     PetscCall(PetscMalloc1(num_coords * dim, &tree->coords_owned));
    265:     PetscCall(PetscArraycpy((PetscReal *)tree->coords_owned, coords, num_coords * dim));
    266:     tree->coords = tree->coords_owned;
    267:     break;
    268:   }

    270:   KDTreeSortContext kd_ctx;
    271:   PetscCall(PetscMalloc2(num_coords * dim, &sorted_indices, num_coords, &temp));
    272:   PetscCall(PetscNew(&kd_ctx));
    273:   kd_ctx->tree = tree;
    274:   for (PetscInt j = 0; j < dim; j++) {
    275:     for (PetscCount i = 0; i < num_coords; i++) sorted_indices[num_coords * j + i] = i;
    276:     kd_ctx->initial_axis = (uint8_t)j;
    277:     PetscCall(PetscTimSort((PetscInt)num_coords, &sorted_indices[num_coords * j], sizeof(*sorted_indices), PetscKDTreeTimSort, kd_ctx));
    278:   }
    279:   PetscCall(PetscFree(kd_ctx));

    281:   PetscInt    num_leaves = (PetscInt)PetscCeilInt64(num_coords, tree->max_bucket_size);
    282:   PetscInt    num_stems  = RoundToNextPowerOfTwo((uint32_t)num_leaves);
    283:   KDTreeBuild kd_build;
    284:   PetscCall(PetscNew(&kd_build));
    285:   kd_build->tree        = tree;
    286:   kd_build->copy_coords = copy_mode == PETSC_COPY_VALUES ? PETSC_TRUE : PETSC_FALSE;
    287:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-kdtree_debug", &kd_build->debug_build, NULL));
    288:   PetscCall(PetscSegBufferCreate(sizeof(KDStem), num_stems, &kd_build->stems));
    289:   PetscCall(PetscSegBufferCreate(sizeof(KDLeaf), num_leaves, &kd_build->leaves));
    290:   PetscCall(PetscSegBufferCreate(sizeof(PetscCount), num_coords, &kd_build->bucket_indices));
    291:   if (kd_build->copy_coords) PetscCall(PetscSegBufferCreate(sizeof(PetscReal), num_coords * dim, &kd_build->bucket_coords));

    293:   PetscCall(PetscKDTreeBuildStemAndLeaves(kd_build, sorted_indices, temp, 0, num_coords, 0, &tree->is_root_leaf, &tree->root_handle));

    295:   PetscCall(PetscSegBufferGetSize(kd_build->stems, &tree->num_stems));
    296:   PetscCall(PetscSegBufferGetSize(kd_build->leaves, &tree->num_leaves));
    297:   PetscCall(PetscSegBufferGetSize(kd_build->bucket_indices, &tree->num_bucket_indices));
    298:   PetscCall(PetscSegBufferExtractAlloc(kd_build->stems, &tree->stems));
    299:   PetscCall(PetscSegBufferExtractAlloc(kd_build->leaves, &tree->leaves));
    300:   PetscCall(PetscSegBufferExtractAlloc(kd_build->bucket_indices, &tree->bucket_indices));
    301:   if (kd_build->copy_coords) {
    302:     PetscCall(PetscFree(tree->coords_owned));
    303:     PetscCall(PetscSegBufferExtractAlloc(kd_build->bucket_coords, &tree->coords_owned));
    304:     tree->coords = tree->coords_owned;
    305:     PetscCall(PetscSegBufferDestroy(&kd_build->bucket_coords));
    306:   }
    307:   PetscCall(PetscSegBufferDestroy(&kd_build->stems));
    308:   PetscCall(PetscSegBufferDestroy(&kd_build->leaves));
    309:   PetscCall(PetscSegBufferDestroy(&kd_build->bucket_indices));
    310:   PetscCall(PetscFree(kd_build));
    311:   PetscCall(PetscFree2(sorted_indices, temp));
    312:   *new_tree = tree;
    313:   PetscCall(PetscLogEventEnd(PetscKDTree_Build, 0, 0, 0, 0));
    314:   PetscFunctionReturn(PETSC_SUCCESS);
    315: }

    317: static inline PetscReal PetscSquareDistance(PetscInt dim, const PetscReal *PETSC_RESTRICT x, const PetscReal *PETSC_RESTRICT y)
    318: {
    319:   PetscReal dist = 0;
    320:   for (PetscInt j = 0; j < dim; j++) dist += PetscSqr(x[j] - y[j]);
    321:   return dist;
    322: }

    324: static inline PetscErrorCode PetscKDTreeQueryLeaf(PetscKDTree tree, KDLeaf leaf, const PetscReal point[], PetscCount *index, PetscReal *distance_sqr)
    325: {
    326:   PetscInt dim = tree->dim;

    328:   PetscFunctionBeginUser;
    329:   *distance_sqr = PETSC_MAX_REAL;
    330:   *index        = -1;
    331:   for (PetscInt i = 0; i < leaf.count; i++) {
    332:     PetscCount point_index = tree->bucket_indices[leaf.indices_handle + i];
    333:     PetscReal  dist        = PetscSquareDistance(dim, point, &tree->coords[point_index * dim]);
    334:     if (dist < *distance_sqr) {
    335:       *distance_sqr = dist;
    336:       *index        = point_index;
    337:     }
    338:   }
    339:   PetscFunctionReturn(PETSC_SUCCESS);
    340: }

    342: static inline PetscErrorCode PetscKDTreeQueryLeaf_CopyCoords(PetscKDTree tree, KDLeaf leaf, const PetscReal point[], PetscCount *index, PetscReal *distance_sqr)
    343: {
    344:   PetscInt dim = tree->dim;

    346:   PetscFunctionBeginUser;
    347:   *distance_sqr = PETSC_MAX_REAL;
    348:   *index        = -1;
    349:   for (PetscInt i = 0; i < leaf.count; i++) {
    350:     // Coord data saved in axis-major ordering for vectorization
    351:     PetscReal dist = 0.;
    352:     for (PetscInt d = 0; d < dim; d++) dist += PetscSqr(point[d] - tree->coords[leaf.coords_handle + d * leaf.count + i]);
    353:     if (dist < *distance_sqr) {
    354:       *distance_sqr = dist;
    355:       *index        = tree->bucket_indices[leaf.indices_handle + i];
    356:     }
    357:   }
    358:   PetscFunctionReturn(PETSC_SUCCESS);
    359: }

    361: // Recursive point query from 'Algorithms for Fast Vector Quantization' by  Sunil Arya and David Mount
    362: // Variant also implemented in pykdtree
    363: 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)
    364: {
    365:   PetscFunctionBeginUser;
    366:   if (*dist_sqr < tol_sqr) PetscFunctionReturn(PETSC_SUCCESS);
    367:   if (is_node_leaf) {
    368:     KDLeaf     leaf = tree->leaves[node_handle];
    369:     PetscReal  dist;
    370:     PetscCount point_index;

    372:     if (leaf.coords_handle > -1) PetscCall(PetscKDTreeQueryLeaf_CopyCoords(tree, leaf, point, &point_index, &dist));
    373:     else PetscCall(PetscKDTreeQueryLeaf(tree, leaf, point, &point_index, &dist));
    374:     if (dist < *dist_sqr) {
    375:       *dist_sqr = dist;
    376:       *index    = point_index;
    377:     }
    378:     PetscFunctionReturn(PETSC_SUCCESS);
    379:   }

    381:   KDStem    stem       = tree->stems[node_handle];
    382:   PetscReal old_offset = offset[stem.axis], new_offset = point[stem.axis] - stem.split;
    383:   if (new_offset <= 0) {
    384:     PetscCall(PetscKDTreeQuery_Recurse(tree, point, stem.less_equal_handle, PetscBTLookup(&stem.are_handles_leaves, LESS_EQUAL_BIT), offset, rd, tol_sqr, index, dist_sqr));
    385:     rd += -PetscSqr(old_offset) + PetscSqr(new_offset);
    386:     if (rd < *dist_sqr) {
    387:       offset[stem.axis] = new_offset;
    388:       PetscCall(PetscKDTreeQuery_Recurse(tree, point, stem.greater_handle, PetscBTLookup(&stem.are_handles_leaves, GREATER_BIT), offset, rd, tol_sqr, index, dist_sqr));
    389:       offset[stem.axis] = old_offset;
    390:     }
    391:   } else {
    392:     PetscCall(PetscKDTreeQuery_Recurse(tree, point, stem.greater_handle, PetscBTLookup(&stem.are_handles_leaves, GREATER_BIT), offset, rd, tol_sqr, index, dist_sqr));
    393:     rd += -PetscSqr(old_offset) + PetscSqr(new_offset);
    394:     if (rd < *dist_sqr) {
    395:       offset[stem.axis] = new_offset;
    396:       PetscCall(PetscKDTreeQuery_Recurse(tree, point, stem.less_equal_handle, PetscBTLookup(&stem.are_handles_leaves, LESS_EQUAL_BIT), offset, rd, tol_sqr, index, dist_sqr));
    397:       offset[stem.axis] = old_offset;
    398:     }
    399:   }
    400:   PetscFunctionReturn(PETSC_SUCCESS);
    401: }

    403: /*@C
    404:   PetscKDTreeQueryPointsNearestNeighbor - find the nearest neighbor in a `PetscKDTree`

    406:   Not Collective, No Fortran Support

    408:   Input Parameters:
    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

    414:   Output Parameters:
    415: + indices   - indices of the nearest neighbor to the query point
    416: - distances - distance between the queried point and the nearest neighbor

    418:   Level: advanced

    420:   Notes:
    421:   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.

    423:   The `indices` and `distances` arrays should be at least of size `num_points`.

    425: .seealso: `PetscKDTree`, `PetscKDTreeCreate()`
    426: @*/
    427: PetscErrorCode PetscKDTreeQueryPointsNearestNeighbor(PetscKDTree tree, PetscCount num_points, const PetscReal points[], PetscReal tolerance, PetscCount indices[], PetscReal distances[])
    428: {
    429:   PetscReal *offsets, rd;

    431:   PetscFunctionBeginUser;
    432:   PetscCall(PetscLogEventBegin(PetscKDTree_Query, 0, 0, 0, 0));
    433:   if (tree == NULL) {
    434:     PetscCheck(num_points == 0, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "num_points may only be zero, if tree is NULL");
    435:     PetscFunctionReturn(PETSC_SUCCESS);
    436:   }
    437:   PetscAssertPointer(points, 3);
    438:   PetscAssertPointer(indices, 5);
    439:   PetscAssertPointer(distances, 6);
    440:   PetscCall(PetscCalloc1(tree->dim, &offsets));

    442:   for (PetscCount p = 0; p < num_points; p++) {
    443:     rd           = 0.;
    444:     distances[p] = PETSC_MAX_REAL;
    445:     indices[p]   = -1;
    446:     PetscCall(PetscKDTreeQuery_Recurse(tree, &points[p * tree->dim], tree->root_handle, (char)tree->is_root_leaf, offsets, rd, PetscSqr(tolerance), &indices[p], &distances[p]));
    447:     distances[p] = PetscSqrtReal(distances[p]);
    448:   }
    449:   PetscCall(PetscFree(offsets));
    450:   PetscCall(PetscLogEventEnd(PetscKDTree_Query, 0, 0, 0, 0));
    451:   PetscFunctionReturn(PETSC_SUCCESS);
    452: }

    454: /*@C
    455:   PetscKDTreeView - view a `PetscKDTree`

    457:   Not Collective, No Fortran Support

    459:   Input Parameters:
    460: + tree   - tree to view
    461: - viewer - visualization context

    463:   Level: advanced

    465: .seealso: `PetscKDTree`, `PetscKDTreeCreate()`, `PetscViewer`
    466: @*/
    467: PetscErrorCode PetscKDTreeView(PetscKDTree tree, PetscViewer viewer)
    468: {
    469:   PetscFunctionBeginUser;
    471:   else PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_SELF, &viewer));
    472:   if (tree == NULL) PetscFunctionReturn(PETSC_SUCCESS);

    474:   PetscCall(PetscViewerASCIIPrintf(viewer, "KDTree:\n"));
    475:   PetscCall(PetscViewerASCIIPushTab(viewer)); // KDTree:
    476:   PetscCall(PetscViewerASCIIPrintf(viewer, "Stems:\n"));
    477:   PetscCall(PetscViewerASCIIPushTab(viewer)); // Stems:
    478:   for (PetscCount i = 0; i < tree->num_stems; i++) {
    479:     KDStem stem = tree->stems[i];
    480:     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",
    481:                                      stem.greater_handle, PetscBTLookup(&stem.are_handles_leaves, LESS_EQUAL_BIT) ? "Leaf" : "Stem", stem.less_equal_handle));
    482:   }
    483:   PetscCall(PetscViewerASCIIPopTab(viewer)); // Stems:

    485:   PetscCall(PetscViewerASCIIPrintf(viewer, "Leaves:\n"));
    486:   PetscCall(PetscViewerASCIIPushTab(viewer)); // Leaves:
    487:   for (PetscCount i = 0; i < tree->num_leaves; i++) {
    488:     KDLeaf leaf = tree->leaves[i];
    489:     PetscCall(PetscViewerASCIIPrintf(viewer, "Leaf %" PetscCount_FMT ": Count=%" PetscInt_FMT, i, leaf.count));
    490:     PetscCall(PetscViewerASCIIPushTab(viewer)); // Coords:
    491:     for (PetscInt j = 0; j < leaf.count; j++) {
    492:       PetscInt   tabs;
    493:       PetscCount bucket_index = tree->bucket_indices[leaf.indices_handle + j];
    494:       PetscCall(PetscViewerASCIIPrintf(viewer, "\n"));
    495:       PetscCall(PetscViewerASCIIPrintf(viewer, "%" PetscCount_FMT ": ", bucket_index));

    497:       PetscCall(PetscViewerASCIIGetTab(viewer, &tabs));
    498:       PetscCall(PetscViewerASCIISetTab(viewer, 0));
    499:       if (leaf.coords_handle > -1) {
    500:         for (PetscInt k = 0; k < tree->dim; k++) PetscCall(PetscViewerASCIIPrintf(viewer, "%g ", (double)tree->coords[leaf.coords_handle + leaf.count * k + j]));
    501:         PetscCall(PetscViewerASCIIPrintf(viewer, " (stored at leaf)"));
    502:       } else {
    503:         for (PetscInt k = 0; k < tree->dim; k++) PetscCall(PetscViewerASCIIPrintf(viewer, "%g ", (double)tree->coords[bucket_index * tree->dim + k]));
    504:       }
    505:       PetscCall(PetscViewerASCIISetTab(viewer, tabs));
    506:     }
    507:     PetscCall(PetscViewerASCIIPopTab(viewer)); // Coords:
    508:     PetscCall(PetscViewerASCIIPrintf(viewer, "\n"));
    509:   }
    510:   PetscCall(PetscViewerASCIIPopTab(viewer)); // Leaves:
    511:   PetscCall(PetscViewerASCIIPopTab(viewer)); // KDTree:
    512:   PetscFunctionReturn(PETSC_SUCCESS);
    513: }