added support for distance fields in stim::kdtree

David Mayerich
1 parent ac48e0d0
Showing 5 changed files with 135 additions and 36 deletions Show diff stats
stim/biomodels/network.h
stim/grids/image_stack.h
stim/math/vector.h
stim/structures/kdtree.cuh
stim/visualization/aabbn.h
@@ -388,6 +388,22 @@ public:
 		return n;
 	}
  
+	//Copy the point cloud representing the centerline for the network into an array
+	void centerline_cloud(T* dst) {
+		size_t p;										//stores the current edge point
+		size_t P;										//stores the number of points in an edge
+		size_t i = 0;									//index into the output array of points
+		for (size_t e = 0; e < E.size(); e++) {			//for each edge in the network
+			P = E[e].size();							//get the number of points in this edge
+			for (p = 0; p < P; p++) {
+				dst[i * 3 + 0] = E[e][p][0];		
+				dst[i * 3 + 1] = E[e][p][1];
+				dst[i * 3 + 2] = E[e][p][2];
+				i++;
+			}
+		}
+	}
+
 	// gaussian function
 	float gaussianFunction(float x, float std=25){ return exp(-x/(2*std*std));} // by default std = 25
  
@@ -418,27 +434,27 @@ public:
 	/// @param sigma is the user-defined tolerance value - smaller values provide a stricter comparison
 	stim::network<T> compare(stim::network<T> A, float sigma, int device){
  
-		stim::network<T> R;								//generate a network storing the result of the comparison
-		R = (*this);									//initialize the result with the current network
+		stim::network<T> R;										//generate a network storing the result of the comparison
+		R = (*this);											//initialize the result with the current network
  
-		T *c;						                 	// centerline (array of double pointers) - points on kdtree must be double
+		T *c;						                 			// centerline (array of double pointers) - points on kdtree must be double
 		size_t n_data = A.total_points();          				// set the number of points
-		c = (T*) malloc(sizeof(T) * n_data * 3); 				
+		c = (T*) malloc(sizeof(T) * n_data * 3);				//allocate an array to store all points in the data set				
  
 		unsigned t = 0;
-		for(unsigned e = 0; e < A.E.size(); e++){					//for each edge in the network
-			for(unsigned p = 0; p < A.E[e].size(); p++){				//for each point in the edge
+		for(unsigned e = 0; e < A.E.size(); e++){				//for each edge in the network
+			for(unsigned p = 0; p < A.E[e].size(); p++){		//for each point in the edge
 				for(unsigned d = 0; d < 3; d++){				//for each coordinate
  
-					c[t * 3 + d] = A.E[e][p][d];
+					c[t * 3 + d] = A.E[e][p][d];				//copy the point into the array c
 				}
 				t++;
 			}
 		}
  
 		//generate a KD-tree for network A
-		//float metric = 0.0;                               				// initialize metric to be returned after comparing the network
-		size_t MaxTreeLevels = 3;									// max tree level
+		//float metric = 0.0;                               	// initialize metric to be returned after comparing the network
+		size_t MaxTreeLevels = 3;								// max tree level
  
 #ifdef __CUDACC__
 		cudaSetDevice(device);
@@ -472,13 +488,13 @@ public:
 		stim::cpu_kdtree<T, 3> kdt;
 		kdt.create(c, n_data, MaxTreeLevels);
 		T *dists = new T[1];								// near neighbor distances
-		size_t *nnIdx = new size_t[1];							// near neighbor indices // allocate near neigh indices
+		size_t *nnIdx = new size_t[1];						// near neighbor indices // allocate near neigh indices
  
 		stim::vec3<T> p0, p1;
 		T m1;
 		T* queryPt = new T[3];
-		for(unsigned e = 0; e < R.E.size(); e++){					//for each edge in A
-			R.E[e].add_mag(0);							//add a new magnitude for the metric
+		for(unsigned e = 0; e < R.E.size(); e++){			//for each edge in A
+			R.E[e].add_mag(0);								//add a new magnitude for the metric
  
 			for(unsigned p = 0; p < R.E[e].size(); p++){				//for each point in the edge
  
@@ -139,7 +139,7 @@ public:
 	/// @param depth,  number of pixels in depth.
 	void init(int channels, int width, int height, int depth)
 	{
-		R.resize(4);
+		//R.resize(4);
 		R[0] = channels;
 		R[1] = width;
 		R[2] = height;
@@ -5,6 +5,7 @@
 #include <cmath>
 #include <sstream>
 #include <vector>
+#include <algorithm>
  
 #include <stim/cuda/cudatools/callable.h>
 #include <stim/math/vec3.h>
@@ -428,12 +428,6 @@ namespace stim {
  
 		search<T, D>(nodes, indices, d_reference_points, d_query_points[idx], &d_indices[idx], &d_distances[idx], idx, next_nodes, next_search_nodes, Judge);    // every query points are independent
 	}
-	template <typename T, int D>
-	__host__ void aaboundingboxing(stim::aabbn<T, D> &bb, T *reference_points, size_t reference_count) {
-		for(size_t i = 0; i < reference_count; i++) {
-			bb.insert(&reference_points[i]);
-		}
-	}
  
 	template <typename T, int D = 3>
 	class cuda_kdtree {
@@ -441,7 +435,7 @@ namespace stim {
 		cuda_kdnode<T> *d_nodes;                                                    																		 
 		size_t *d_index;
 		kdtree::point<T, D>* d_reference_points;
-		size_t d_reference_count;
+		size_t npts;
 		int num_nodes;
 	public:
 		~cuda_kdtree() {
@@ -449,13 +443,18 @@ namespace stim {
 			HANDLE_ERROR(cudaFree(d_index));
 			HANDLE_ERROR(cudaFree(d_reference_points));
 		}
-		void create(T *h_reference_points, size_t reference_count, size_t max_levels, stim::aabbn<T, D> &bb) {
+
+		/// 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) {
 			if (max_levels > 10) {
 				std::cout<<"The max_tree_levels should be smaller!"<<std::endl;
 				exit(1);
 			}		
-			bb.init(&h_reference_points[0]);
-			aaboundingboxing<T, D>(bb, h_reference_points, reference_count);
+			//bb.init(&h_reference_points[0]);
+			//aaboundingboxing<T, D>(bb, h_reference_points, reference_count);
  
 			std::vector <kdtree::point<T, D>> reference_points(reference_count);																				// restore the reference points in particular way
 			for (size_t j = 0; j < reference_count; j++)
@@ -465,14 +464,14 @@ namespace stim {
 			tree.cpu_create(reference_points, max_levels);																											// building a tree on cpu
 			kdtree::kdnode<T> *d_root = tree.get_root();
 			num_nodes = tree.get_num_nodes();
-			d_reference_count = reference_count;																												// also equals to reference_count
+			npts = reference_count;																												// also equals to reference_count
  
 			HANDLE_ERROR(cudaMalloc((void**)&d_nodes, sizeof(cuda_kdnode<T>) * num_nodes));																		// copy data from host to device
-			HANDLE_ERROR(cudaMalloc((void**)&d_index, sizeof(size_t) * d_reference_count));
-			HANDLE_ERROR(cudaMalloc((void**)&d_reference_points, sizeof(kdtree::point<T, D>) * d_reference_count));
+			HANDLE_ERROR(cudaMalloc((void**)&d_index, sizeof(size_t) * npts));
+			HANDLE_ERROR(cudaMalloc((void**)&d_reference_points, sizeof(kdtree::point<T, D>) * npts));
  
 			std::vector < cuda_kdnode<T> > tmp_nodes(num_nodes);																									
-			std::vector <size_t> indices(d_reference_count);
+			std::vector <size_t> indices(npts);
 			std::vector <kdtree::kdnode<T>*> next_nodes;
 			size_t cur_pos = 0;
 			next_nodes.push_back(d_root);
@@ -511,6 +510,12 @@ namespace stim {
 			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<T, D>) * reference_points.size(), cudaMemcpyHostToDevice));
 		}
+
+		/// 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<T, D> > query_points(query_count);
 			for (size_t j = 0; j < query_count; j++)
@@ -536,7 +541,7 @@ namespace stim {
 			HANDLE_ERROR(cudaMalloc((void**)&next_search_nodes, threads * blocks * 50 * sizeof(int)));	
 			HANDLE_ERROR(cudaMemcpy(d_query_points, &query_points[0], sizeof(T) * query_points.size() * D, cudaMemcpyHostToDevice));
  
-			search_batch<<<threads, blocks>>> (d_nodes, d_index, d_reference_points, d_query_points, query_points.size(), d_indices, d_distances, next_nodes, next_search_nodes, Judge);
+			search_batch<<<blocks, threads>>> (d_nodes, d_index, d_reference_points, d_query_points, query_points.size(), d_indices, d_distances, next_nodes, next_search_nodes, Judge);
  
 			if (Judge == NULL) {																																// do the following work if the thread works safely
 				HANDLE_ERROR(cudaMemcpy(indices, d_indices, sizeof(size_t) * query_points.size(), cudaMemcpyDeviceToHost));
@@ -549,6 +554,63 @@ 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<T, D> getbox() {
+			size_t N = npts;
+			//std::vector < typename kdtree::point<T, D> > 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
+
+			stim::aabbn<T, D> bb(cpu_ref);
+
+			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<T, 3> 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<<"     "<<Q[i * 3 + 0] << "     " << Q[i * 3 + 1] << "     " << Q[i * 3 + 2] << std::endl;
+					}
+				}
+			}
+			size_t* temp = (size_t*)malloc(N * sizeof(size_t));	//allocate space to store the indices (unused)
+			search(Q, N, temp, dist);
+		}
+
+		//generate an implicit distance field for the KD-tree
+		void dist_field3(T* dist, size_t* dims) {
+			stim::aabbn<T, D> bb = getbox();					//get a bounding box around the tree
+			dist_field3(dist, dims, bb);
+		}
+
 	};
 }				//end namespace stim
 #endif
 \ No newline at end of file
@@ -15,17 +15,22 @@ struct aabbn{
 	T low[D];					//top left corner position
 	T high[D]; 							//dimensions along x and y and z
  
-//public:
-
-	CUDA_CALLABLE aabbn(){					//initialize an axis aligned bounding box of size 0 at the given position
-		for(size_t d = 0; d < D; d++)
-			low[d] = high[d] = 0;
+	CUDA_CALLABLE void init(T* i) {
+		for (size_t d = 0; d < D; d++)
+			low[d] = high[d] = i[d];
 	}
  
-	CUDA_CALLABLE void init(T* init){
-		for(size_t d = 0; d < D; d++)
-			low[d] = high[d] = init[d];
+//public:
+
+	//CUDA_CALLABLE aabbn(){					//initialize an axis aligned bounding box of size 0 at the given position
+	//	for(size_t d = 0; d < D; d++)
+	//		low[d] = high[d] = 0;
+	//}
+	CUDA_CALLABLE aabbn() {}
+	CUDA_CALLABLE aabbn(T* i) {
+		init(i);
 	}
+	
  
 	//insert a point into the bounding box, growing the box appropriately
 	CUDA_CALLABLE void insert(T* p){
@@ -46,6 +51,21 @@ struct aabbn{
 			if(low[d] > b[d]) low[d] = b[d];
 	}
  
+	CUDA_CALLABLE T length(size_t d) {
+		return high[d] - low[d];
+	}
+
+	CUDA_CALLABLE aabbn<T, D> operator*(T s) {
+		aabbn<T, D> newbox;
+		for (size_t d = 0; d < D; d++) {
+			T c = (low[d] + high[d]) / 2;
+			T l = high[d] - low[d];
+			newbox.low[d] = c - l * s / 2;
+			newbox.high[d] = c + l * s / 2;
+		}
+		return newbox;
+	}
+
 };
  
 }