[WIP] KDTree implementation

Verifying the recursive building on several input data.

kdtree/main.pyx

@@ -67,7 +67,7 @@ cdef extern from *:
                 I split_dim,
                 I split_index,
                 I n_features,
-                I n_points) except +
+                I n_points) nogil except +
 
 
 cdef I_t partition_node_indices(
@@ -76,7 +76,7 @@ cdef I_t partition_node_indices(
         I_t split_dim,
         I_t split_index,
         I_t n_features,
-        I_t n_points) except -1:
+        I_t n_points) nogil except -1:
     """Partition points in the node into two equal-sized groups.
     Upon return, the values in node_indices will be rearranged such that
     (assuming numpy-style indexing):
@@ -126,7 +126,8 @@ cdef I_t partition_node_indices(
 cdef cypclass Node activable:
     """A KDTree Node"""
 
-    double[:] point
+    D_t * point
+    I_t n_dims
     active Node left
     active Node right
 
@@ -143,30 +144,37 @@ cdef cypclass Node activable:
             I_t n,
             I_t depth,
             I_t n_dims,
-            I_t dim
+            I_t dim,
+            I_t start,
+            I_t end,
     ):
         cdef I_t i
-        if (depth < 0):
+        cdef I_t nn = end - start
+        cdef I_t split_index = (start + end) // 2
+
+        self.n_dims = n_dims
+
+        if (depth < 0) or (nn <= 0):
             return
 
-        printf("Depth %d\n", depth)
-        printf("Number of dimensions %d\n", n_dims)
-        printf("Splitting dimension %d\n", dim)
-        for i in range(n):
-            printf("X[%d, %d] = %f\n",
-                    i, dim,
-                    points[i*n_dims + dim])
+        printf("Depth %d\nSplitting dimension "
+        "%d\nSplit indice %d\nSorting in [%d, %d]\n\n\n", depth, dim,
+        split_index, start, end)
+
+        partition_node_indices(points + start,
+                indices + start, dim, nn // 2, n_dims, nn)
+        self.point = points + split_index
 
-        printf("\n")
-        # TODO: find a way to partitions indices here
 
         self.left = activate(consume Node())
         self.right = activate(consume Node())
 
+        next_dim = (dim + 1) % n_dims
+
         self.left.build_node(NULL, points, indices,
-                n, depth - 1, n_dims, (dim + 1) % n_dims)
+                n, depth - 1, n_dims, next_dim, start, split_index)
         self.right.build_node(NULL, points, indices,
-                n, depth - 1, n_dims, (dim + 1) % n_dims)
+                n, depth - 1, n_dims, next_dim, split_index, end)
 
 
 cdef I_t start() nogil:
@@ -175,7 +183,7 @@ cdef I_t start() nogil:
     cdef I_t i
     cdef I_t n = 12
     cdef I_t d = 2
-    cdef I_t depth = 5
+    cdef I_t depth = 10
 
     # TODO: use memory view for convenience
     # cdef D_t p[12][2]
@@ -188,8 +196,8 @@ cdef I_t start() nogil:
 
     for i in range(n):
         indices[i] = i
-        points[i * d] = (i / golden_ratio) % 1
-        points[i * d + 1] = i / n
+        points[i * d] = i # (i / golden_ratio) % 1
+        points[i * d + 1] = i # i / n
 
     printf("Before\n")
     node = consume Node()
@@ -197,10 +205,15 @@ cdef I_t start() nogil:
         return -1
 
     root = activate(consume node)
-    root.build_node(NULL, points, indices, n, depth, d, 0)
-    scheduler.finish()
+    root.build_node(NULL, points, indices, n,
+            depth, d, dim=0, start=0, end=n)
 
+    scheduler.finish()
     del scheduler
+
+    for i in range(n):
+        printf("indices[%d] = %d\n", i, indices[i])
+
     free(points)
     free(indices)