1 static const char help[] = "Test KDTree\n\n"; 2 3 #include <petsc.h> 4 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 } 11 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; 22 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; 27 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 *)©_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(); 41 42 coords_size = num_coords * dim; 43 PetscCall(PetscMalloc1(coords_size, &coords)); 44 PetscCall(PetscRandomCreate(comm, &random)); 45 46 for (PetscInt loop_count = 0; loop_count < loops; loop_count++) { 47 PetscCall(PetscRandomGetValuesReal(random, coords_size, coords)); 48 49 PetscCall(PetscKDTreeCreate(num_coords, dim, coords, copy_mode, bucket_size, &tree)); 50 num_trees_built++; 51 if (view_tree) PetscCall(PetscKDTreeView(tree, NULL)); 52 53 if (test_tree_points) { // Test querying the current coordinates 54 PetscCount *indices; 55 PetscReal *distances; 56 num_points_queried += num_coords; 57 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 } 65 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; 71 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)); 76 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 } 95 96 if (view_performance) { 97 PetscLogEvent kdtree_build_log, kdtree_query_log; 98 PetscEventPerfInfo build_perf_info, query_perf_info; 99 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 } 109 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 } 116 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 133 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*/ 145