Actual source code: ex1.c

  1: static const char help[] = "Test KDTree\n\n";

  3: #include <petsc.h>

  5: static inline PetscReal Distance(PetscInt dim, const PetscReal *PETSC_RESTRICT x, const PetscReal *PETSC_RESTRICT y)
  6: {
  7:   PetscReal dist = 0;
  8:   for (PetscInt j = 0; j < dim; j++) dist += PetscSqr(x[j] - y[j]);
  9:   return PetscSqrtReal(dist);
 10: }

 12: int main(int argc, char **argv)
 13: {
 14:   MPI_Comm      comm;
 15:   PetscInt      num_coords, dim, num_rand_points = 0, bucket_size = PETSC_DECIDE;
 16:   PetscRandom   random;
 17:   PetscReal    *coords;
 18:   PetscInt      coords_size, num_points_queried = 0, num_trees_built = 0, loops = 1;
 19:   PetscBool     view_tree = PETSC_FALSE, view_performance = PETSC_TRUE, test_tree_points = PETSC_FALSE, test_rand_points = PETSC_FALSE, query_rand_points = PETSC_FALSE;
 20:   PetscCopyMode copy_mode = PETSC_OWN_POINTER;
 21:   PetscKDTree   tree;

 23:   PetscFunctionBeginUser;
 24:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 25:   PetscCall(PetscLogDefaultBegin()); // So we can use PetscLogEventGetPerfInfo without -log_view
 26:   comm = PETSC_COMM_WORLD;

 28:   PetscOptionsBegin(PETSC_COMM_WORLD, "", "Test Options", "none");
 29:   PetscCall(PetscOptionsInt("-num_coords", "Number of coordinates", "", PETSC_FALSE, &num_coords, NULL));
 30:   PetscCall(PetscOptionsInt("-dim", "Dimension of the coordinates", "", PETSC_FALSE, &dim, NULL));
 31:   PetscCall(PetscOptionsInt("-bucket_size", "Size of the bucket at each leaf", "", bucket_size, &bucket_size, NULL));
 32:   PetscCall(PetscOptionsBoundedInt("-loops", "Number of times to loop through KDTree creation and querying", "", loops, &loops, NULL, 1));
 33:   PetscCall(PetscOptionsEnum("-copy_mode", "Copy mode to use with KDTree", "PetscKDTreeCreate", PetscCopyModes, (PetscEnum)copy_mode, (PetscEnum *)&copy_mode, NULL));
 34:   PetscCall(PetscOptionsBool("-view_tree", "View the KDTree", "", view_tree, &view_tree, NULL));
 35:   PetscCall(PetscOptionsBool("-view_performance", "View the performance speed of the KDTree", "", view_performance, &view_performance, NULL));
 36:   PetscCall(PetscOptionsBool("-test_tree_points", "Test querying tree points against itself", "", test_tree_points, &test_tree_points, NULL));
 37:   PetscCall(PetscOptionsBool("-test_rand_points", "Test querying random points via brute force", "", test_rand_points, &test_rand_points, NULL));
 38:   PetscCall(PetscOptionsBool("-query_rand_points", "Query querying random points", "", query_rand_points, &query_rand_points, NULL));
 39:   if (test_rand_points || query_rand_points) PetscCall(PetscOptionsInt("-num_rand_points", "Number of random points to test with", "", num_rand_points, &num_rand_points, NULL));
 40:   PetscOptionsEnd();

 42:   coords_size = num_coords * dim;
 43:   PetscCall(PetscMalloc1(coords_size, &coords));
 44:   PetscCall(PetscRandomCreate(comm, &random));

 46:   for (PetscInt loop_count = 0; loop_count < loops; loop_count++) {
 47:     PetscCall(PetscRandomGetValuesReal(random, coords_size, coords));

 49:     PetscCall(PetscKDTreeCreate(num_coords, dim, coords, copy_mode, bucket_size, &tree));
 50:     num_trees_built++;
 51:     if (view_tree) PetscCall(PetscKDTreeView(tree, NULL));

 53:     if (test_tree_points) { // Test querying the current coordinates
 54:       PetscCount *indices;
 55:       PetscReal  *distances;
 56:       num_points_queried += num_coords;

 58:       PetscCall(PetscMalloc2(num_coords, &indices, num_coords, &distances));
 59:       PetscCall(PetscKDTreeQueryPointsNearestNeighbor(tree, num_coords, coords, PETSC_MACHINE_EPSILON * 1e3, indices, distances));
 60:       for (PetscInt i = 0; i < num_coords; i++) {
 61:         if (indices[i] != i) PetscCall(PetscPrintf(comm, "Query failed for coord %" PetscInt_FMT ", query returned index %" PetscCount_FMT " with distance %g\n", i, indices[i], (double)distances[i]));
 62:       }
 63:       PetscCall(PetscFree2(indices, distances));
 64:     }

 66:     if (num_rand_points > 0) {
 67:       PetscCount *indices;
 68:       PetscReal  *distances;
 69:       PetscReal  *rand_points;
 70:       PetscInt    rand_queries_size = num_rand_points * dim;

 72:       num_points_queried += num_rand_points;
 73:       PetscCall(PetscMalloc3(rand_queries_size, &rand_points, num_rand_points, &indices, num_rand_points, &distances));
 74:       PetscCall(PetscRandomGetValuesReal(random, rand_queries_size, rand_points));
 75:       PetscCall(PetscKDTreeQueryPointsNearestNeighbor(tree, num_rand_points, rand_points, 0., indices, distances));

 77:       if (test_rand_points) {
 78:         for (PetscInt i = 0; i < num_rand_points; i++) {
 79:           PetscReal *rand_point = &rand_points[dim * i], nearest_distance = PETSC_MAX_REAL;
 80:           PetscInt   index = -1;
 81:           for (PetscInt j = 0; j < num_coords; j++) {
 82:             PetscReal dist = Distance(dim, rand_point, &coords[dim * j]);
 83:             if (dist < nearest_distance) {
 84:               nearest_distance = dist;
 85:               index            = j;
 86:             }
 87:           }
 88:           if (indices[i] != index)
 89:             PetscCall(PetscPrintf(comm, "Query failed for random point %" PetscInt_FMT ". Query returned index %" PetscCount_FMT " with distance %g, but coordinate %" PetscInt_FMT " with distance %g is closer\n", i, indices[i], (double)distances[i], index, (double)nearest_distance));
 90:         }
 91:       }
 92:       PetscCall(PetscFree3(rand_points, indices, distances));
 93:     }
 94:   }

 96:   if (view_performance) {
 97:     PetscLogEvent      kdtree_build_log, kdtree_query_log;
 98:     PetscEventPerfInfo build_perf_info, query_perf_info;

100:     PetscCall(PetscLogEventGetId("KDTreeBuild", &kdtree_build_log));
101:     PetscCall(PetscLogEventGetPerfInfo(PETSC_DETERMINE, kdtree_build_log, &build_perf_info));
102:     PetscCall(PetscLogEventGetId("KDTreeQuery", &kdtree_query_log));
103:     PetscCall(PetscLogEventGetPerfInfo(PETSC_DETERMINE, kdtree_query_log, &query_perf_info));
104:     PetscCall(PetscPrintf(comm, "KDTreeBuild %.6e for %" PetscInt_FMT " trees built\n", build_perf_info.time, num_trees_built));
105:     PetscCall(PetscPrintf(comm, "\tTime per tree: %.6e\n", build_perf_info.time / num_trees_built));
106:     PetscCall(PetscPrintf(comm, "KDTreeQuery %.6e for %" PetscCount_FMT " queries\n", query_perf_info.time, (PetscCount)num_points_queried));
107:     PetscCall(PetscPrintf(comm, "\tTime per query: %.6e\n", query_perf_info.time / num_points_queried));
108:   }

110:   if (copy_mode != PETSC_OWN_POINTER) PetscCall(PetscFree(coords));
111:   PetscCall(PetscKDTreeDestroy(&tree));
112:   PetscCall(PetscRandomDestroy(&random));
113:   PetscCall(PetscFinalize());
114:   return 0;
115: }

117: /*TEST
118:   testset:
119:     suffix: kdtree
120:     args: -num_coords 35 -test_tree_points -test_rand_points -num_rand_points 300 -bucket_size 13 -view_performance false -view_tree true -kdtree_debug
121:     test:
122:       suffix: 1D
123:       args: -dim 1
124:     test:
125:       suffix: 2D
126:       args: -dim 2
127:     test:
128:       suffix: 3D
129:       args: -dim 3
130:     test:
131:       suffix: 3D_small
132:       args: -dim 3 -num_coords 13

134:   testset:
135:     suffix: kdtree_3D_large
136:     args: -dim 3 -num_coords 350 -test_tree_points -test_rand_points -num_rand_points 300 -view_performance false -kdtree_debug
137:     test:
138:     test:
139:       suffix: copy
140:       args: -copy_mode copy_values
141:     test:
142:       suffix: use
143:       args: -copy_mode use_pointer
144: TEST*/