diff --git a/stim/structures/kdtree.cuh b/stim/structures/kdtree.cuh
index c6b350e..50cece8 100644
--- a/stim/structures/kdtree.cuh
+++ b/stim/structures/kdtree.cuh
@@ -18,6 +18,7 @@
 #include <iostream>
 #include <algorithm>
 #include <stim/cuda/cudatools/error.h>
+#include <stim/visualization/aabbn.h>
 
 namespace stim {
 	namespace kdtree {
@@ -427,6 +428,12 @@ 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 {
@@ -435,17 +442,21 @@ namespace stim {
 		size_t *d_index;
 		kdtree::point<T, D>* d_reference_points;
 		size_t d_reference_count;
+		int num_nodes;
 	public:
 		~cuda_kdtree() {
 			HANDLE_ERROR(cudaFree(d_nodes));
 			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) {
+		void create(T *h_reference_points, size_t reference_count, size_t max_levels, stim::aabbn<T, D> &bb) {
 			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);
+
 			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++)
 				for (size_t i = 0; i < D; i++)
@@ -453,7 +464,7 @@ namespace stim {
 			cpu_kdtree<T, D> tree;																																// creating a tree on cpu
 			tree.cpu_create(reference_points, max_levels);																											// building a tree on cpu
 			kdtree::kdnode<T> *d_root = tree.get_root();
-			int num_nodes = tree.get_num_nodes();
+			num_nodes = tree.get_num_nodes();
 			d_reference_count = 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
--
libgit2 0.21.4