Move point creation to the front-end

kdtree/kdtree.pyx

@@ -54,7 +54,7 @@ cdef extern from *:
     }
     """
     void partition_node_indices[D, I](
-                D *data,
+                D *data_ptr,
                 I *node_indices,
                 I start,
                 I mid,
@@ -79,8 +79,8 @@ cdef cypclass Node activable:
 
     void build_node(
             self,
-            D_t * data,
-            I_t * indices,
+            D_t * data_ptr,
+            I_t * indices_ptr,
             I_t depth,
             I_t n_dims,
             I_t dim,
@@ -96,18 +96,18 @@ cdef cypclass Node activable:
         if (depth < 0) or (end - start <= 1):
             return
 
-        partition_node_indices(data, indices, start, mid, end, dim, n_dims)
-        self.point = data + mid
+        partition_node_indices(data_ptr, indices_ptr, start, mid, end, dim, n_dims)
+        self.point = data_ptr + mid
 
         self.left = consume Node()
         self.right = consume Node()
 
         self.left.build_node(NULL,
-                             data, indices,
+                             data_ptr, indices_ptr,
                              depth - 1, n_dims, next_dim,
                              start, mid)
         self.right.build_node(NULL,
-                              data, indices,
+                              data_ptr, indices_ptr,
                               depth - 1, n_dims, next_dim,
                               mid, end)
 
@@ -131,12 +131,10 @@ cdef cypclass KDTree:
     np.ndarray data_arr
     np.ndarray idx_array_arr
 
-
-
     active Node root
 
-    D_t *data
-    I_t *indices
+    D_t *data_ptr
+    I_t *indices_ptr
 
     __init__(self,
              np.ndarray data,
@@ -150,23 +148,15 @@ cdef cypclass KDTree:
         self.d = d
         self.depth = depth
 
-        # TODO: define it on the front-end
-        # Use Golden Spiral for the layout
-        cdef D_t golden_ratio = (1 + 5**0.5)/2
-        self.data = <D_t *> malloc(n * d * sizeof(D_t))
-        self.indices = <I_t *> malloc(n * sizeof(I_t))
-
+        self.data_ptr = <D_t *> data.data
+        self.indices_ptr = <I_t *> malloc(n * sizeof(I_t))
         for i in range(n):
-            self.indices[i] = i
-            self.data[i * d] = (i / golden_ratio) % 1
-            self.data[i * d + 1] = i / n
-
-        # TODO: end
+            self.indices_ptr[i] = i
 
         self._recursive_build()
 
     void _recursive_build(self):
-        # TODO: introducing a context for the runtime
+        # TODO: introducing a context manager for the runtime
         # would be nice here:
         #   ```
         #   with scheduler:
@@ -184,22 +174,23 @@ cdef cypclass KDTree:
         # a new argument is prepredend for the Promise,
         # which we aren't using using here, hence the extra NULL.
         self.root.build_node(NULL,
-                             self.data,
-                             self.indices,
+                             self.data_ptr,
+                             self.indices_ptr,
                              depth, n_dims=d, dim=0, start=0, end=n)
 
         scheduler.finish()
         del scheduler
 
         for i in range(n):
-            printf("X[indices[%d] = %d]= %f %f\n", i, self.indices[i]
-                   self.data[self.indices[i] * d],
-                   self.data[self.indices[i] * d + 1])
+            printf("X[indices[%d] = %d]= %f %f\n",
+                   i, self.indices_ptr[i],
+                   self.data_ptr[self.indices_ptr[i] * d],
+                   self.data_ptr[self.indices_ptr[i] * d + 1])
 
     void __dealloc__(self):
         printf("Deallocating KDTree datastructures\n")
-        free(self.data)
-        free(self.indices)
+        free(self.data_ptr)
+        free(self.indices_ptr)
         printf("Done deallocating KDTree datastructures\n")
         # XXX It sometimes blocks over here and then segfaults
 

kdtree/main.py

@@ -2,5 +2,10 @@ import numpy as np
 import kdtree
 
 if __name__ == "__main__":
-    X = np.random.randint(0, 100, size=(24, 2))
+    n, d = 24, 2
+    golden_ratio = (1 + 5 ** 0.5) / 2
+    X = np.zeros((n, d))
+    for i in range(n):
+        X[i, 0] = (i / golden_ratio) % 1
+        X[i, 1] = i / n
     tree = kdtree.KDTree(X, depth=10)