Commits

Matthew Turk committed 6e93560

Initial import of Quadtree example code

Comments (0)

Files changed (2)

+cimport numpy as np
+import numpy as np
+from libc.stdlib cimport malloc, free
+
+DEF NREF=128
+DEF MAXLEVEL=40
+
+cdef struct qnode:
+    # Note we do not store any center or width info
+    # That is only held during the traversal
+    qnode **children
+    np.int64_t *indices
+    np.float64_t *pos[2]
+    np.int32_t n
+
+cdef class QuadTree:
+    cdef np.float64_t left_edge[2]
+    cdef np.float64_t right_edge[2]
+    cdef qnode* base
+
+    def __init__(self, left_edge, right_edge):
+        cdef int i
+        for i in range(2):
+            self.left_edge[i] = left_edge[i]
+            self.right_edge[i] = right_edge[i]
+        self.base = create_node()
+
+    cdef np.int64_t find_nearest(self, np.float64_t pos[2]):
+        cdef qnode* node = self.find(pos)
+        # We have our node, now we need the nearest neighbor
+        cdef int i, mi
+        cdef np.float64_t min_r2 = 1e300
+        cdef np.float64_t radius2 = 0.0
+        mi = -1
+        for i in range(node.n):
+            radius2  = (node.pos[0][i] - pos[0])**2.0
+            radius2 += (node.pos[1][i] - pos[1])**2.0
+            if radius2 < min_r2:
+                mi = i
+                min_r2 = radius2
+        if mi == -1:
+            print "COULD NOT FIND ANYTHING", min_r2, radius2, node.n
+            raise RuntimeError
+        return node.indices[mi]
+
+    cdef qnode *find(self, np.float64_t pos[2]):
+        cdef np.float64_t c[2], w[2]
+        cdef int i, level, ci
+        cdef qnode *node = self.base
+        for i in range(2):
+            c[i] = (self.right_edge[i] + self.left_edge[i])/2.0
+            w[i] = (self.right_edge[i] - self.left_edge[i])
+        level = 0
+        check_refine_node(node, c, w)
+        while node.children != NULL:
+            level += 1
+            if level > MAXLEVEL: raise RuntimeError
+            ci = 0
+            if pos[0] > c[0]:
+                ci += 2
+                c[0] += w[0] / 4.0
+            else:
+                c[0] -= w[0] / 4.0
+            if pos[1] > c[1]:
+                ci += 1
+                c[1] += w[1] / 4.0
+            else:
+                c[1] -= w[1] / 4.0
+            w[0] /= 2.0
+            w[1] /= 2.0
+            node = node.children[ci]
+            check_refine_node(node, c, w)
+        return node
+
+    def add(self, np.ndarray[np.float64_t, ndim=2] positions):
+        cdef int n, i
+        cdef np.float64_t pos[2]
+        cdef qnode *node
+        for i in range(positions.shape[0]):
+            pos[0] = positions[i, 0]
+            pos[1] = positions[i, 1]
+            node = self.find(pos)
+            node.pos[0][node.n] = pos[0]
+            node.pos[1][node.n] = pos[1]
+            node.indices[node.n] = i
+            node.n += 1
+
+    def __dealloc__(self):
+        clear_node(self.base)
+
+    def nodes(self):
+        cdef int n
+        
+    def print_stats(self):
+        ls = self.stats()
+        for i in range(MAXLEVEL):
+            if ls[i] == 0: break
+            print "% 3i  % 10i" % (i, ls[i])
+
+    def stats(self):
+        cdef np.ndarray[np.int64_t, ndim=1] level_counts
+        level_counts = np.zeros(MAXLEVEL, dtype="int64")
+        cdef int level = 0
+        self.recursively_count(self.base, level_counts, 0)
+        return level_counts
+
+    cdef void recursively_count(self, qnode *node,
+                                np.ndarray[np.int64_t, ndim=1] counts,
+                                int level):
+        cdef int i
+        if level > MAXLEVEL: return
+        counts[level] += 1
+        if node.children == NULL:
+            return
+        for i in range(4):
+            self.recursively_count(node.children[i], counts, level + 1)
+
+    def interp(self,
+               np.ndarray[np.float64_t, ndim=1] LE,
+               np.ndarray[np.float64_t, ndim=1] RE,
+               np.ndarray[np.float64_t, ndim=2] dest,
+               np.ndarray[np.float64_t, ndim=1] values,
+               np.float64_t max_r = -1.0):
+        # First we get the bounds of our grid in x, y plane.
+        cdef np.float64_t pos[2]
+        cdef qnode *node
+        cdef int i, j, k
+        cdef np.float64_t dds[2], r
+        if max_r > 0:
+            max_r = max_r * max_r
+        cdef np.int64_t ind
+        for i in range(2):
+            dds[i] = (RE[i]-LE[i])/dest.shape[i]
+        for i in range(dest.shape[0]):
+            pos[0] = LE[0] + (i + 0.5) * dds[0]
+            for j in range(dest.shape[1]):
+                pos[1] = LE[1] + (j + 0.5) * dds[1]
+                node = self.find(pos)
+                if max_r > 0:
+                    r = pos[0]*pos[0] + pos[1]*pos[1]
+                    if r > max_r: continue
+                if node.n == 0:
+                    continue
+                ind = self.find_nearest(pos)
+                dest[i,j] = values[ind]
+
+cdef qnode* create_node():
+    cdef int i
+    cdef qnode* nn = <qnode*> malloc(sizeof(qnode))
+    nn.children = NULL
+    nn.n = 0
+    nn.indices = <np.int64_t*> malloc(sizeof(np.int64_t) * NREF)
+    nn.pos[0] = <np.float64_t*> malloc(sizeof(np.float64_t) * NREF)
+    nn.pos[1] = <np.float64_t*> malloc(sizeof(np.float64_t) * NREF)
+    for i in range(NREF):
+        nn.indices[i] = -1
+    return nn
+
+cdef void clear_node(qnode *tc):
+    cdef int i
+    if tc == NULL: return
+    if tc.children != NULL:
+        for i in range(4):
+            clear_node(tc.children[i])
+        free(tc.children) # Free the pointer
+    else:
+        free(tc.pos[0])
+        free(tc.pos[1])
+        free(tc.indices)
+    free(tc)
+
+cdef void check_refine_node(qnode *node, np.float64_t c[2],
+                            np.float64_t w[2]):
+    cdef int i, ci
+    cdef np.float64_t w2[2], c2[2]
+    cdef qnode *child
+    if node.n < NREF: return
+    # We need to refine this node
+    node.children = <qnode **> malloc(sizeof(qnode*) * 4)
+    for i in range(4): # Four new ones
+        node.children[i] = create_node()
+    for i in range(NREF):
+        # ci == 0: left of both
+        # ci == 1: left of c[0], right of c[0]
+        # ci == 2: right of c[0], left of c[0]
+        # ci == 3: right of both
+        ci = 0
+        if node.pos[0][i] > c[0]:
+            ci += 2
+        if node.pos[1][i] > c[1]:
+            ci += 1
+        child = node.children[ci]
+        child.pos[0][child.n] = node.pos[0][i]
+        child.pos[1][child.n] = node.pos[1][i]
+        child.indices[child.n] = node.indices[i]
+        child.n += 1
+    free(node.indices)
+    free(node.pos[0])
+    free(node.pos[1])
+    node.indices = node.pos[0] = node.pos[1] = NULL
+    node.n = 0
+    w2[0] = w[0]/2.0
+    w2[1] = w[1]/2.0
+    for i in range(4):
+        if i < 2:
+            c2[0] = c[0] - w[0]/4.0
+        else:
+            c2[0] = c[0] + w[0]/4.0
+        if i == 0 or i == 2:
+            c2[1] = c[1] - w[1]/4.0
+        else:
+            c2[1] = c[1] + w[1]/4.0
+        check_refine_node(node.children[i], c2, w2)
+
+import numpy as np
+import matplotlib.pyplot as pyplot
+import pyximport
+pyximport.install(setup_args={'include_dirs':[np.get_include()]})
+
+# We generate two different clouds with a Gaussian distribution
+cloud1 = 0.2*np.random.randn( 1000,2) + 0.25
+cloud2 = 0.1*np.random.randn(10000,2) + 0.75
+
+# Let's come up with some data for these clouds -- their radius to the cloud
+# center.
+r1 = ((cloud1 - 0.25)**2).sum(axis=1)**0.5
+r2 = ((cloud2 - 0.75)**2).sum(axis=1)**0.5
+
+# Now we can plot them, individually as different colors, so that we can see
+# what they look like.
+pyplot.scatter(cloud1[:,0], cloud1[:,1], c='b', marker='.')
+pyplot.scatter(cloud2[:,0], cloud2[:,1], c='r', marker='.')
+pyplot.axis("equal")
+pyplot.xlim(-0.5, 1.5)
+pyplot.ylim(-0.5, 1.5)
+pyplot.savefig("clouds.png")
+
+# Now we load our QuadTree Cython module.
+import quadtree
+clouds = np.concatenate([cloud1, cloud2])
+# We need bounds, so we simply calculate the left and right edges of our
+# points.
+LE = clouds.min(axis=0)
+RE = clouds.max(axis=0)
+QT = quadtree.QuadTree(LE, RE)
+QT.add(clouds)
+QT.print_stats()
+
+# Concatenate our values, which will be looked-up into
+values = np.concatenate([r1, r2])
+# We create an empty canvas of 4096^2 pixels
+dest = np.zeros((4096, 4096))
+QT.interp(LE, RE, dest, values)
+
+pyplot.clf()
+pyplot.imshow(dest, interpolation='nearest', origin='lower',
+              extent = (-0.5, 1.5, -0.5, 1.5))
+pyplot.savefig("r_interp.png")