1 static const char help[] = "Test KDTree\n\n";
2
3 #include <petsc.h>
4
Distance(PetscInt dim,const PetscReal * PETSC_RESTRICT x,const PetscReal * PETSC_RESTRICT y)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
main(int argc,char ** argv)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 output_file: output/empty.out
138 test:
139 test:
140 suffix: copy
141 args: -copy_mode copy_values
142 test:
143 suffix: use
144 args: -copy_mode use_pointer
145 TEST*/
146