diff --git a/stim/structures/kdtree.cuh b/stim/structures/kdtree.cuh index 43dbd2d..093c71a 100644 --- a/stim/structures/kdtree.cuh +++ b/stim/structures/kdtree.cuh @@ -1,4 +1,4 @@ -// right now the size of CUDA STACK is set to 1000, increase it if you mean to make deeper tree +// right now the size of CUDA STACK is set to 50, increase it if you mean to make deeper tree // data should be stored in row-major // x1,x2,x3,x4,x5...... // y1,y2,y3,y4,y5...... @@ -22,16 +22,16 @@ #include namespace stim { - namespace kdtree { + namespace cpu_kdtree { template // typename refers to float or double while D refers to dimension of points struct point { T dim[D]; // create a structure to store every one input point }; template - class kdnode { + class cpu_kdnode { public: - kdnode() { // constructor for initializing a kdnode + cpu_kdnode() { // constructor for initializing a kdnode parent = NULL; // set every node's parent, left and right kdnode pointers to NULL left = NULL; right = NULL; @@ -42,258 +42,12 @@ namespace stim { } int idx; // index of current node int parent_idx, left_idx, right_idx; // index of parent, left and right nodes - kdnode *parent, *left, *right; // parent, left and right kdnodes + cpu_kdnode *parent, *left, *right; // parent, left and right kdnodes T split_value; // splitting value of current node std::vector indices; // it indicates the points' indices that current node has size_t level; // tree level of current node }; - } // end of namespace kdtree - - template // set dimension of data to default 3 - class cpu_kdtree { - protected: - int current_axis; // current judging axis - int n_id; // store the total number of nodes - std::vector < typename kdtree::point > *tmp_points; // transfer or temperary points - std::vector < typename kdtree::point > cpu_tmp_points; // for cpu searching - kdtree::kdnode *root; // root node - static cpu_kdtree *cur_tree_ptr; - public: - cpu_kdtree() { // constructor for creating a cpu_kdtree - cur_tree_ptr = this; // create a class pointer points to the current class value - n_id = 0; // set total number of points to default 0 - } - - ~cpu_kdtree() { // destructor of cpu_kdtree - std::vector *> next_nodes; - next_nodes.push_back(root); - while (next_nodes.size()) { - std::vector *> next_search_nodes; - while (next_nodes.size()) { - kdtree::kdnode *cur = next_nodes.back(); - next_nodes.pop_back(); - if (cur->left) - next_search_nodes.push_back(cur->left); - if (cur->right) - next_search_nodes.push_back(cur->right); - delete cur; - } - next_nodes = next_search_nodes; - } - root = NULL; - } - - void cpu_create(std::vector < typename kdtree::point > &reference_points, size_t max_levels) { - tmp_points = &reference_points; - root = new kdtree::kdnode(); // initializing the root node - root->idx = n_id++; // the index of root is 0 - root->level = 0; // tree level begins at 0 - root->indices.resize(reference_points.size()); // get the number of points - for (size_t i = 0; i < reference_points.size(); i++) { - root->indices[i] = i; // set indices of input points - } - std::vector *> next_nodes; // next nodes - next_nodes.push_back(root); // push back the root node - while (next_nodes.size()) { - std::vector *> next_search_nodes; // next search nodes - while (next_nodes.size()) { // two same WHILE is because we need to make a new vector to store nodes for search - kdtree::kdnode *current_node = next_nodes.back(); // handle node one by one (right first) - next_nodes.pop_back(); // pop out current node in order to store next round of nodes - if (current_node->level < max_levels) { - if (current_node->indices.size() > 1) { // split if the nonleaf node contains more than one point - kdtree::kdnode *left = new kdtree::kdnode(); - kdtree::kdnode *right = new kdtree::kdnode(); - left->idx = n_id++; // set the index of current node's left node - right->idx = n_id++; - split(current_node, left, right); // split left and right and determine a node - std::vector temp; // empty vecters of int - //temp.resize(current_node->indices.size()); - current_node->indices.swap(temp); // clean up current node's indices - current_node->left = left; - current_node->right = right; - current_node->left_idx = left->idx; - current_node->right_idx = right->idx; - if (right->indices.size()) - next_search_nodes.push_back(right); // left pop out first - if (left->indices.size()) - next_search_nodes.push_back(left); - } - } - } - next_nodes = next_search_nodes; // go deeper within the tree - } - } - - static bool sort_points(const size_t a, const size_t b) { // create functor for std::sort - std::vector < typename kdtree::point > &pts = *cur_tree_ptr->tmp_points; // put cur_tree_ptr to current input points' pointer - return pts[a].dim[cur_tree_ptr->current_axis] < pts[b].dim[cur_tree_ptr->current_axis]; - } - - void split(kdtree::kdnode *cur, kdtree::kdnode *left, kdtree::kdnode *right) { - std::vector < typename kdtree::point > &pts = *tmp_points; - current_axis = cur->level % D; // indicate the judicative dimension or axis - std::sort(cur->indices.begin(), cur->indices.end(), sort_points); // using SortPoints as comparison function to sort the data - size_t mid_value = cur->indices[cur->indices.size() / 2]; // odd in the mid_value, even take the floor - cur->split_value = pts[mid_value].dim[current_axis]; // get the parent node - left->parent = cur; // set the parent of the next search nodes to current node - right->parent = cur; - left->level = cur->level + 1; // level + 1 - right->level = cur->level + 1; - left->parent_idx = cur->idx; // set its parent node's index - right->parent_idx = cur->idx; - for (size_t i = 0; i < cur->indices.size(); i++) { // split into left and right half-space one by one - size_t idx = cur->indices[i]; - if (pts[idx].dim[current_axis] < cur->split_value) - left->indices.push_back(idx); - else - right->indices.push_back(idx); - } - } - - void create(T *h_reference_points, size_t reference_count, size_t max_levels) { - std::vector < typename kdtree::point > reference_points(reference_count); // restore the reference points in particular way - for (size_t j = 0; j < reference_count; j++) - for (size_t i = 0; i < D; i++) - reference_points[j].dim[i] = h_reference_points[j * D + i]; - cpu_create(reference_points, max_levels); - cpu_tmp_points = *tmp_points; - } - - int get_num_nodes() const { // get the total number of nodes - return n_id; - } - - kdtree::kdnode* get_root() const { // get the root node of tree - return root; - } - - T cpu_distance(const kdtree::point &a, const kdtree::point &b) { - T distance = 0; - - for (size_t i = 0; i < D; i++) { - T d = a.dim[i] - b.dim[i]; - distance += d*d; - } - return distance; - } - - void cpu_search_at_node(kdtree::kdnode *cur, const kdtree::point &query, size_t *index, T *distance, kdtree::kdnode **node) { - T best_distance = FLT_MAX; // initialize the best distance to max of floating point - size_t best_index = 0; - std::vector < typename kdtree::point > pts = cpu_tmp_points; - while (true) { - size_t split_axis = cur->level % D; - if (cur->left == NULL) { // risky but acceptable, same goes for right because left and right are in same pace - *node = cur; // pointer points to a pointer - for (size_t i = 0; i < cur->indices.size(); i++) { - size_t idx = cur->indices[i]; - T d = cpu_distance(query, pts[idx]); // compute distances - /// if we want to compute k nearest neighbor, we can input the last resul - /// (last_best_dist < dist < best_dist) to select the next point until reaching to k - if (d < best_distance) { - best_distance = d; - best_index = idx; // record the nearest neighbor index - } - } - break; // find the target point then break the loop - } - else if (query.dim[split_axis] < cur->split_value) { // if it has son node, visit the next node on either left side or right side - cur = cur->left; - } - else { - cur = cur->right; - } - } - *index = best_index; - *distance = best_distance; - } - - void cpu_search_at_node_range(kdtree::kdnode *cur, const kdtree::point &query, T range, size_t *index, T *distance) { - T best_distance = FLT_MAX; // initialize the best distance to max of floating point - size_t best_index = 0; - std::vector < typename kdtree::point > pts = cpu_tmp_points; - std::vector < typename kdtree::kdnode*> next_node; - next_node.push_back(cur); - while (next_node.size()) { - std::vector*> next_search; - while (next_node.size()) { - cur = next_node.back(); - next_node.pop_back(); - size_t split_axis = cur->level % D; - if (cur->left == NULL) { - for (size_t i = 0; i < cur->indices.size(); i++) { - size_t idx = cur->indices[i]; - T d = cpu_distance(query, pts[idx]); - if (d < best_distance) { - best_distance = d; - best_index = idx; - } - } - } - else { - T d = query.dim[split_axis] - cur->split_value; // computer distance along specific axis or dimension - /// there are three possibilities: on either left or right, and on both left and right - if (fabs(d) > range) { // absolute value of floating point to see if distance will be larger that best_dist - if (d < 0) - next_search.push_back(cur->left); // every left[split_axis] is less and equal to cur->split_value, so it is possible to find the nearest point in this region - else - next_search.push_back(cur->right); - } - else { // it is possible that nereast neighbor will appear on both left and right - next_search.push_back(cur->left); - next_search.push_back(cur->right); - } - } - } - next_node = next_search; // pop out at least one time - } - *index = best_index; - *distance = best_distance; - } - - void cpu_search(T *h_query_points, size_t query_count, size_t *h_indices, T *h_distances) { - /// first convert the input query point into specific type - kdtree::point query; - for (size_t j = 0; j < query_count; j++) { - for (size_t i = 0; i < D; i++) - query.dim[i] = h_query_points[j * D + i]; - /// find the nearest node, this will be the upper bound for the next time searching - kdtree::kdnode *best_node = NULL; - T best_distance = FLT_MAX; - size_t best_index = 0; - T radius = 0; // radius for range - cpu_search_at_node(root, query, &best_index, &best_distance, &best_node); // simple search to rougly determine a result for next search step - radius = sqrt(best_distance); // It is possible that nearest will appear in another region - /// find other possibilities - kdtree::kdnode *cur = best_node; - while (cur->parent != NULL) { // every node that you pass will be possible to be the best node - /// go up - kdtree::kdnode *parent = cur->parent; // travel back to every node that we pass through - size_t split_axis = (parent->level) % D; - /// search other nodes - size_t tmp_index; - T tmp_distance = FLT_MAX; - if (fabs(parent->split_value - query.dim[split_axis]) <= radius) { - /// search opposite node - if (parent->left != cur) - cpu_search_at_node_range(parent->left, query, radius, &tmp_index, &tmp_distance); // to see whether it is its mother node's left son node - else - cpu_search_at_node_range(parent->right, query, radius, &tmp_index, &tmp_distance); - } - if (tmp_distance < best_distance) { - best_distance = tmp_distance; - best_index = tmp_index; - } - cur = parent; - } - h_indices[j] = best_index; - h_distances[j] = best_distance; - } - } - }; //end class kdtree - - template - cpu_kdtree* cpu_kdtree::cur_tree_ptr = NULL; // definition of cur_tree_ptr pointer points to the current class + } // end of namespace cpu_kdtree template struct cuda_kdnode { @@ -305,7 +59,7 @@ namespace stim { }; template - __device__ T gpu_distance(kdtree::point &a, kdtree::point &b) { + __device__ T gpu_distance(cpu_kdtree::point &a, cpu_kdtree::point &b) { T distance = 0; for (size_t i = 0; i < D; i++) { @@ -316,7 +70,7 @@ namespace stim { } template - __device__ void search_at_node(cuda_kdnode *nodes, size_t *indices, kdtree::point *d_reference_points, int cur, kdtree::point &d_query_point, size_t *d_index, T *d_distance, int *d_node) { + __device__ void search_at_node(cuda_kdnode *nodes, size_t *indices, cpu_kdtree::point *d_reference_points, int cur, cpu_kdtree::point &d_query_point, size_t *d_index, T *d_distance, int *d_node) { T best_distance = FLT_MAX; size_t best_index = 0; @@ -346,7 +100,7 @@ namespace stim { } template - __device__ void search_at_node_range(cuda_kdnode *nodes, size_t *indices, kdtree::point *d_reference_points, kdtree::point &d_query_point, int cur, T range, size_t *d_index, T *d_distance, size_t id, int *next_nodes, int *next_search_nodes, int *Judge) { + __device__ void search_at_node_range(cuda_kdnode *nodes, size_t *indices, cpu_kdtree::point *d_reference_points, cpu_kdtree::point &d_query_point, int cur, T range, size_t *d_index, T *d_distance, size_t id, int *next_nodes, int *next_search_nodes, int *Judge) { T best_distance = FLT_MAX; size_t best_index = 0; @@ -405,7 +159,7 @@ namespace stim { } template - __device__ void search(cuda_kdnode *nodes, size_t *indices, kdtree::point *d_reference_points, kdtree::point &d_query_point, size_t *d_index, T *d_distance, size_t id, int *next_nodes, int *next_search_nodes, int *Judge) { + __device__ void search(cuda_kdnode *nodes, size_t *indices, cpu_kdtree::point *d_reference_points, cpu_kdtree::point &d_query_point, size_t *d_index, T *d_distance, size_t id, int *next_nodes, int *next_search_nodes, int *Judge) { int best_node = 0; T best_distance = FLT_MAX; size_t best_index = 0; @@ -438,7 +192,7 @@ namespace stim { } template - __global__ void search_batch(cuda_kdnode *nodes, size_t *indices, kdtree::point *d_reference_points, kdtree::point *d_query_points, size_t d_query_count, size_t *d_indices, T *d_distances, int *next_nodes, int *next_search_nodes, int *Judge) { + __global__ void search_batch(cuda_kdnode *nodes, size_t *indices, cpu_kdtree::point *d_reference_points, cpu_kdtree::point *d_query_points, size_t d_query_count, size_t *d_indices, T *d_distances, int *next_nodes, int *next_search_nodes, int *Judge) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= d_query_count) return; // avoid segfault @@ -446,11 +200,11 @@ namespace stim { } template - void search_stream(cuda_kdnode *d_nodes, size_t *d_index, kdtree::point *d_reference_points, kdtree::point *query_stream_points, size_t stream_count, size_t *indices, T *distances) { + void search_stream(cuda_kdnode *d_nodes, size_t *d_index, cpu_kdtree::point *d_reference_points, cpu_kdtree::point *query_stream_points, size_t stream_count, size_t *indices, T *distances) { unsigned int threads = (unsigned int)(stream_count > 1024 ? 1024 : stream_count); unsigned int blocks = (unsigned int)(stream_count / threads + (stream_count % threads ? 1 : 0)); - kdtree::point *d_query_points; + cpu_kdtree::point *d_query_points; size_t *d_indices; T *d_distances; @@ -480,26 +234,121 @@ namespace stim { HANDLE_ERROR(cudaFree(d_distances)); } - template - class cuda_kdtree { + template // set dimension of data to default 3 + class kdtree { protected: - cuda_kdnode *d_nodes; - size_t *d_index; - kdtree::point* d_reference_points; - size_t npts; - int num_nodes; + int current_axis; // current judging axis + int n_id; // store the total number of nodes + std::vector < typename cpu_kdtree::point > *tmp_points; // transfer or temperary points + std::vector < typename cpu_kdtree::point > cpu_tmp_points; // for cpu searching + cpu_kdtree::cpu_kdnode *root; // root node + static kdtree *cur_tree_ptr; + #ifdef __CUDACC__ + cuda_kdnode *d_nodes; + size_t *d_index; + cpu_kdtree::point* d_reference_points; + size_t npts; + int num_nodes; + #endif public: - ~cuda_kdtree() { + kdtree() { // constructor for creating a cpu_kdtree + cur_tree_ptr = this; // create a class pointer points to the current class value + n_id = 0; // set total number of points to default 0 + } + + ~kdtree() { // destructor of cpu_kdtree + std::vector *> next_nodes; + next_nodes.push_back(root); + while (next_nodes.size()) { + std::vector *> next_search_nodes; + while (next_nodes.size()) { + cpu_kdtree::cpu_kdnode *cur = next_nodes.back(); + next_nodes.pop_back(); + if (cur->left) + next_search_nodes.push_back(cur->left); + if (cur->right) + next_search_nodes.push_back(cur->right); + delete cur; + } + next_nodes = next_search_nodes; + } + root = NULL; + #ifdef __CUDACC__ HANDLE_ERROR(cudaFree(d_nodes)); HANDLE_ERROR(cudaFree(d_index)); HANDLE_ERROR(cudaFree(d_reference_points)); + #endif } - - /// Create a KD-tree given a pointer to an array of reference points and the number of reference points - /// @param h_reference_points is a host array containing the reference points in (x0, y0, z0, ...., ) order - /// @param reference_count is the number of reference point in the array - /// @param max_levels is the deepest number of tree levels allowed - void create(T *h_reference_points, size_t reference_count, size_t max_levels = 3) { + + void cpu_create(std::vector < typename cpu_kdtree::point > &reference_points, size_t max_levels) { + tmp_points = &reference_points; + root = new cpu_kdtree::cpu_kdnode(); // initializing the root node + root->idx = n_id++; // the index of root is 0 + root->level = 0; // tree level begins at 0 + root->indices.resize(reference_points.size()); // get the number of points + for (size_t i = 0; i < reference_points.size(); i++) { + root->indices[i] = i; // set indices of input points + } + std::vector *> next_nodes; // next nodes + next_nodes.push_back(root); // push back the root node + while (next_nodes.size()) { + std::vector *> next_search_nodes; // next search nodes + while (next_nodes.size()) { // two same WHILE is because we need to make a new vector to store nodes for search + cpu_kdtree::cpu_kdnode *current_node = next_nodes.back(); // handle node one by one (right first) + next_nodes.pop_back(); // pop out current node in order to store next round of nodes + if (current_node->level < max_levels) { + if (current_node->indices.size() > 1) { // split if the nonleaf node contains more than one point + cpu_kdtree::cpu_kdnode *left = new cpu_kdtree::cpu_kdnode(); + cpu_kdtree::cpu_kdnode *right = new cpu_kdtree::cpu_kdnode(); + left->idx = n_id++; // set the index of current node's left node + right->idx = n_id++; + split(current_node, left, right); // split left and right and determine a node + std::vector temp; // empty vecters of int + //temp.resize(current_node->indices.size()); + current_node->indices.swap(temp); // clean up current node's indices + current_node->left = left; + current_node->right = right; + current_node->left_idx = left->idx; + current_node->right_idx = right->idx; + if (right->indices.size()) + next_search_nodes.push_back(right); // left pop out first + if (left->indices.size()) + next_search_nodes.push_back(left); + } + } + } + next_nodes = next_search_nodes; // go deeper within the tree + } + } + + static bool sort_points(const size_t a, const size_t b) { // create functor for std::sort + std::vector < typename cpu_kdtree::point > &pts = *cur_tree_ptr->tmp_points; // put cur_tree_ptr to current input points' pointer + return pts[a].dim[cur_tree_ptr->current_axis] < pts[b].dim[cur_tree_ptr->current_axis]; + } + + void split(cpu_kdtree::cpu_kdnode *cur, cpu_kdtree::cpu_kdnode *left, cpu_kdtree::cpu_kdnode *right) { + std::vector < typename cpu_kdtree::point > &pts = *tmp_points; + current_axis = cur->level % D; // indicate the judicative dimension or axis + std::sort(cur->indices.begin(), cur->indices.end(), sort_points); // using SortPoints as comparison function to sort the data + size_t mid_value = cur->indices[cur->indices.size() / 2]; // odd in the mid_value, even take the floor + cur->split_value = pts[mid_value].dim[current_axis]; // get the parent node + left->parent = cur; // set the parent of the next search nodes to current node + right->parent = cur; + left->level = cur->level + 1; // level + 1 + right->level = cur->level + 1; + left->parent_idx = cur->idx; // set its parent node's index + right->parent_idx = cur->idx; + for (size_t i = 0; i < cur->indices.size(); i++) { // split into left and right half-space one by one + size_t idx = cur->indices[i]; + if (pts[idx].dim[current_axis] < cur->split_value) + left->indices.push_back(idx); + else + right->indices.push_back(idx); + } + } + + void create(T *h_reference_points, size_t reference_count, size_t max_levels) { + #ifdef __CUDACC__ if (max_levels > 10) { std::cout<<"The max_tree_levels should be smaller!"<(bb, h_reference_points, reference_count); - std::vector < typename kdtree::point> reference_points(reference_count); // restore the reference points in particular way + std::vector < typename cpu_kdtree::point> reference_points(reference_count); // restore the reference points in particular way for (size_t j = 0; j < reference_count; j++) for (size_t i = 0; i < D; i++) - reference_points[j].dim[i] = h_reference_points[j * D + i]; - cpu_kdtree tree; // creating a tree on cpu - tree.cpu_create(reference_points, max_levels); // building a tree on cpu - kdtree::kdnode *d_root = tree.get_root(); - num_nodes = tree.get_num_nodes(); + reference_points[j].dim[i] = h_reference_points[j * D + i]; // creating a tree on cpu + (*this).cpu_create(reference_points, max_levels); // building a tree on cpu + cpu_kdtree::cpu_kdnode *d_root = (*this).get_root(); + num_nodes = (*this).get_num_nodes(); npts = reference_count; // also equals to reference_count HANDLE_ERROR(cudaMalloc((void**)&d_nodes, sizeof(cuda_kdnode) * num_nodes)); // copy data from host to device HANDLE_ERROR(cudaMalloc((void**)&d_index, sizeof(size_t) * npts)); - HANDLE_ERROR(cudaMalloc((void**)&d_reference_points, sizeof(kdtree::point) * npts)); + HANDLE_ERROR(cudaMalloc((void**)&d_reference_points, sizeof(cpu_kdtree::point) * npts)); std::vector < cuda_kdnode > tmp_nodes(num_nodes); std::vector indices(npts); - std::vector *> next_nodes; + std::vector *> next_nodes; size_t cur_pos = 0; next_nodes.push_back(d_root); while (next_nodes.size()) { - std::vector *> next_search_nodes; + std::vector *> next_search_nodes; while (next_nodes.size()) { - kdtree::kdnode *cur = next_nodes.back(); + cpu_kdtree::cpu_kdnode *cur = next_nodes.back(); next_nodes.pop_back(); int id = cur->idx; // the nodes at same level are independent tmp_nodes[id].level = cur->level; @@ -559,16 +407,154 @@ namespace stim { } HANDLE_ERROR(cudaMemcpy(d_nodes, &tmp_nodes[0], sizeof(cuda_kdnode) * tmp_nodes.size(), cudaMemcpyHostToDevice)); HANDLE_ERROR(cudaMemcpy(d_index, &indices[0], sizeof(size_t) * indices.size(), cudaMemcpyHostToDevice)); - HANDLE_ERROR(cudaMemcpy(d_reference_points, &reference_points[0], sizeof(kdtree::point) * reference_count, cudaMemcpyHostToDevice)); + HANDLE_ERROR(cudaMemcpy(d_reference_points, &reference_points[0], sizeof(cpu_kdtree::point) * reference_count, cudaMemcpyHostToDevice)); + + #else + std::vector < typename cpu_kdtree::point > reference_points(reference_count); // restore the reference points in particular way + for (size_t j = 0; j < reference_count; j++) + for (size_t i = 0; i < D; i++) + reference_points[j].dim[i] = h_reference_points[j * D + i]; + cpu_create(reference_points, max_levels); + cpu_tmp_points = *tmp_points; + + #endif + } + + int get_num_nodes() const { // get the total number of nodes + return n_id; + } + + cpu_kdtree::cpu_kdnode* get_root() const { // get the root node of tree + return root; + } + + T cpu_distance(const cpu_kdtree::point &a, const cpu_kdtree::point &b) { + T distance = 0; + + for (size_t i = 0; i < D; i++) { + T d = a.dim[i] - b.dim[i]; + distance += d*d; + } + return distance; + } + + void cpu_search_at_node(cpu_kdtree::cpu_kdnode *cur, const cpu_kdtree::point &query, size_t *index, T *distance, cpu_kdtree::cpu_kdnode **node) { + T best_distance = FLT_MAX; // initialize the best distance to max of floating point + size_t best_index = 0; + std::vector < typename cpu_kdtree::point > pts = cpu_tmp_points; + while (true) { + size_t split_axis = cur->level % D; + if (cur->left == NULL) { // risky but acceptable, same goes for right because left and right are in same pace + *node = cur; // pointer points to a pointer + for (size_t i = 0; i < cur->indices.size(); i++) { + size_t idx = cur->indices[i]; + T d = cpu_distance(query, pts[idx]); // compute distances + /// if we want to compute k nearest neighbor, we can input the last resul + /// (last_best_dist < dist < best_dist) to select the next point until reaching to k + if (d < best_distance) { + best_distance = d; + best_index = idx; // record the nearest neighbor index + } + } + break; // find the target point then break the loop + } + else if (query.dim[split_axis] < cur->split_value) { // if it has son node, visit the next node on either left side or right side + cur = cur->left; + } + else { + cur = cur->right; + } + } + *index = best_index; + *distance = best_distance; + } + + void cpu_search_at_node_range(cpu_kdtree::cpu_kdnode *cur, const cpu_kdtree::point &query, T range, size_t *index, T *distance) { + T best_distance = FLT_MAX; // initialize the best distance to max of floating point + size_t best_index = 0; + std::vector < typename cpu_kdtree::point > pts = cpu_tmp_points; + std::vector < typename cpu_kdtree::cpu_kdnode*> next_node; + next_node.push_back(cur); + while (next_node.size()) { + std::vector*> next_search; + while (next_node.size()) { + cur = next_node.back(); + next_node.pop_back(); + size_t split_axis = cur->level % D; + if (cur->left == NULL) { + for (size_t i = 0; i < cur->indices.size(); i++) { + size_t idx = cur->indices[i]; + T d = cpu_distance(query, pts[idx]); + if (d < best_distance) { + best_distance = d; + best_index = idx; + } + } + } + else { + T d = query.dim[split_axis] - cur->split_value; // computer distance along specific axis or dimension + /// there are three possibilities: on either left or right, and on both left and right + if (fabs(d) > range) { // absolute value of floating point to see if distance will be larger that best_dist + if (d < 0) + next_search.push_back(cur->left); // every left[split_axis] is less and equal to cur->split_value, so it is possible to find the nearest point in this region + else + next_search.push_back(cur->right); + } + else { // it is possible that nereast neighbor will appear on both left and right + next_search.push_back(cur->left); + next_search.push_back(cur->right); + } + } + } + next_node = next_search; // pop out at least one time + } + *index = best_index; + *distance = best_distance; + } + + void cpu_search(T *h_query_points, size_t query_count, size_t *h_indices, T *h_distances) { + /// first convert the input query point into specific type + cpu_kdtree::point query; + for (size_t j = 0; j < query_count; j++) { + for (size_t i = 0; i < D; i++) + query.dim[i] = h_query_points[j * D + i]; + /// find the nearest node, this will be the upper bound for the next time searching + cpu_kdtree::cpu_kdnode *best_node = NULL; + T best_distance = FLT_MAX; + size_t best_index = 0; + T radius = 0; // radius for range + cpu_search_at_node(root, query, &best_index, &best_distance, &best_node); // simple search to rougly determine a result for next search step + radius = sqrt(best_distance); // It is possible that nearest will appear in another region + /// find other possibilities + cpu_kdtree::cpu_kdnode *cur = best_node; + while (cur->parent != NULL) { // every node that you pass will be possible to be the best node + /// go up + cpu_kdtree::cpu_kdnode *parent = cur->parent; // travel back to every node that we pass through + size_t split_axis = (parent->level) % D; + /// search other nodes + size_t tmp_index; + T tmp_distance = FLT_MAX; + if (fabs(parent->split_value - query.dim[split_axis]) <= radius) { + /// search opposite node + if (parent->left != cur) + cpu_search_at_node_range(parent->left, query, radius, &tmp_index, &tmp_distance); // to see whether it is its mother node's left son node + else + cpu_search_at_node_range(parent->right, query, radius, &tmp_index, &tmp_distance); + } + if (tmp_distance < best_distance) { + best_distance = tmp_distance; + best_index = tmp_index; + } + cur = parent; + } + h_indices[j] = best_index; + h_distances[j] = best_distance; + } } - /// Search the KD tree for nearest neighbors to a set of specified query points - /// @param h_query_points an array of query points in (x0, y0, z0, ...) order - /// @param query_count is the number of query points - /// @param indices are the indices to the nearest reference point for each query points - /// @param distances is an array containing the distance between each query point and the nearest reference point void search(T *h_query_points, size_t query_count, size_t *indices, T *distances) { - std::vector < typename kdtree::point > query_points(query_count); + #ifdef __CUDACC__ + std::vector < typename cpu_kdtree::point > query_points(query_count); for (size_t j = 0; j < query_count; j++) for (size_t i = 0; i < D; i++) query_points[j].dim[i] = h_query_points[j * D + i]; @@ -595,7 +581,7 @@ namespace stim { unsigned int threads = (unsigned int)(query_count > 1024 ? 1024 : query_count); unsigned int blocks = (unsigned int)(query_count / threads + (query_count % threads ? 1 : 0)); - kdtree::point *d_query_points; // create a pointer pointing to query points on gpu + cpu_kdtree::point *d_query_points; // create a pointer pointing to query points on gpu size_t *d_indices; T *d_distances; @@ -624,64 +610,18 @@ namespace stim { HANDLE_ERROR(cudaFree(d_indices)); HANDLE_ERROR(cudaFree(d_distances)); } - } - - /// Return the number of points in the KD tree - size_t num_points() { - return npts; - } - stim::aabbn getbox() { - size_t N = npts; - //std::vector < typename kdtree::point > cpu_ref(npts); //allocate space on the CPU for the reference points - T* cpu_ref = (T*)malloc(N * D * sizeof(T)); //allocate space on the CPU for the reference points - HANDLE_ERROR(cudaMemcpy(cpu_ref, d_reference_points, N * D * sizeof(T), cudaMemcpyDeviceToHost)); //copy from GPU to CPU + #else + cpu_search(h_query_points, query_count, indices, distances); - stim::aabbn bb(cpu_ref); + #endif - for (size_t i = 1; i < N; i++) { //for each reference point - //std::cout << "( " << cpu_ref[i * D + 0] << ", " << cpu_ref[i * D + 1] << ", " << cpu_ref[i * D + 2] << ")" << std::endl; - bb.insert(&cpu_ref[i * D]); - } - return bb; } - //generate an implicit distance field for the KD-tree - void dist_field3(T* dist, size_t* dims, stim::aabbn bb) { - size_t N = 1; //number of query points that make up the distance field - for (size_t d = 0; d < 3; d++) N *= dims[d]; //calculate the total number of query points - - //calculate the grid spatial parameters - T dx = 0; - if (dims[0] > 1) dx = bb.length(0) / dims[0]; - T dy = 0; - if (dims[1] > 1) dy = bb.length(1) / dims[1]; - T dz = 0; - if (dims[2] > 1) dz = bb.length(2) / dims[2]; - - T* Q = (T*)malloc(N * 3 * sizeof(T)); //allocate space for the query points - size_t i; - for (size_t z = 0; z < dims[2]; z++) { //for each query point (which is a point in the grid) - for (size_t y = 0; y < dims[1]; y++) { - for (size_t x = 0; x < dims[0]; x++) { - i = z * dims[1] * dims[0] + y * dims[0] + x; - Q[i * 3 + 0] = bb.low[0] + x * dx + dx / 2; - Q[i * 3 + 1] = bb.low[1] + y * dy + dy / 2; - Q[i * 3 + 2] = bb.low[2] + z * dz + dz / 2; - //std::cout << i<<" "< bb = getbox(); //get a bounding box around the tree - dist_field3(dist, dims, bb); - } + template + kdtree* kdtree::cur_tree_ptr = NULL; // definition of cur_tree_ptr pointer points to the current class - }; } //end namespace stim #endif \ No newline at end of file diff --git a/stim/visualization/gl_network.h b/stim/visualization/gl_network.h index 694efe7..28704a9 100644 --- a/stim/visualization/gl_network.h +++ b/stim/visualization/gl_network.h @@ -44,6 +44,23 @@ public: } /// Render the network centerline as a series of line strips. + /// glCenterline0 is for only one input + void glCenterline0(){ + if (!glIsList(dlist)) { //if dlist isn't a display list, create it + dlist = glGenLists(1); //generate a display list + glNewList(dlist, GL_COMPILE); //start a new display list + for (unsigned e = 0; e < E.size(); e++) { //for each edge in the network + glBegin(GL_LINE_STRIP); + for (unsigned p = 0; p < E[e].size(); p++) { //for each point on that edge + glVertex3f(E[e][p][0], E[e][p][1], E[e][p][2]); //set the vertex position based on the current point + glTexCoord1f(0); //set white color + } + glEnd(); + } + glEndList(); //end the display list + } + glCallList(dlist); // render the display list + } /// @param m specifies the magnitude value used as the vertex weight (radius, error, etc.) void glCenterline(unsigned m = 0){ -- libgit2 0.21.4