[Scipy-svn] r4777 - in branches/spatial/scipy/spatial: . tests

scipy-svn@scip... scipy-svn@scip...
Sun Oct 5 00:52:04 CDT 2008


Author: peridot
Date: 2008-10-05 00:51:47 -0500 (Sun, 05 Oct 2008)
New Revision: 4777

Modified:
   branches/spatial/scipy/spatial/kdtree.py
   branches/spatial/scipy/spatial/tests/test_kdtree.py
Log:
Added all-neighbors code, both array-of-points and between two kd-trees. Not optimized.


Modified: branches/spatial/scipy/spatial/kdtree.py
===================================================================
--- branches/spatial/scipy/spatial/kdtree.py	2008-10-05 05:02:32 UTC (rev 4776)
+++ branches/spatial/scipy/spatial/kdtree.py	2008-10-05 05:51:47 UTC (rev 4777)
@@ -4,6 +4,14 @@
 from heapq import heappush, heappop
 
 def distance_p(x,y,p=2):
+    """Compute the pth power of the L**p distance between x and y
+    
+    For efficiency, this function computes the L**p distance but does
+    not extract the pth root. If p is 1 or infinity, this is equal to
+    the actual L**p distance.
+    """
+    x = np.asarray(x)
+    y = np.asarray(y)
     if p==np.inf:
         return np.amax(np.abs(y-x),axis=-1)
     elif p==1:
@@ -11,12 +19,59 @@
     else:
         return np.sum(np.abs(y-x)**p,axis=-1)
 def distance(x,y,p=2):
+    """Compute the L**p distance between x and y"""
+    x = np.asarray(x)
+    y = np.asarray(y)
     if p==np.inf or p==1:
         return distance_p(x,y,p)
     else:
         return distance_p(x,y,p)**(1./p)
 
+class Rectangle(object):
+    """Hyperrectangle class.
 
+    Represents a Cartesian product of intervals.
+    """
+    def __init__(self, maxes, mins):
+        """Construct a hyperrectangle."""
+        self.maxes = np.maximum(maxes,mins).astype(np.float)
+        self.mins = np.minimum(maxes,mins).astype(np.float)
+        self.k, = self.maxes.shape
+
+    def __repr__(self):
+        return "<Rectangle %s>" % zip(self.mins, self.maxes)
+
+    def volume(self):
+        """Total volume."""
+        return np.prod(self.maxes-self.mins)
+    
+    def split(self, d, split):
+        """Produce two hyperrectangles by splitting along axis d."""
+        mid = np.copy(self.maxes)
+        mid[d] = split
+        less = Rectangle(self.mins, mid)
+        mid = np.copy(self.mins)
+        mid[d] = split
+        greater = Rectangle(mid, self.maxes)
+        return less, greater
+
+    def min_distance_point(self, x, p=2.):
+        """Compute the minimum distance between x and a point in the hyperrectangle."""
+        return distance(0, np.maximum(0,np.maximum(self.mins-x,x-self.maxes)),p)
+
+    def max_distance_point(self, x, p=2.):
+        """Compute the maximum distance between x and a point in the hyperrectangle."""
+        return distance(0, np.maximum(self.maxes-x,x-self.mins),p)
+
+    def min_distance_rectangle(self, other, p=2.):
+        """Compute the minimum distance between points in the two hyperrectangles."""
+        return distance(0, np.maximum(0,np.maximum(self.mins-other.maxes,other.mins-self.maxes)),p)
+
+    def max_distance_rectangle(self, other, p=2.):
+        """Compute the maximum distance between points in the two hyperrectangles."""
+        return distance(0, np.maximum(self.maxes-other.mins,other.maxes-self.mins),p)
+
+
 class KDTree(object):
     """kd-tree for quick nearest-neighbor lookup
 
@@ -42,6 +97,11 @@
     For large dimensions (20 is already large) do not expect this to run 
     significantly faster than brute force. High-dimensional nearest-neighbor
     queries are a substantial open problem in computer science.
+
+    The tree also supports all-neighbors queries, both with arrays of points
+    and with other kd-trees. These do use a reasonably efficient algorithm,
+    but the kd-tree is not necessarily the best data structure for this
+    sort of calculation.
     """
 
     def __init__(self, data, leafsize=10):
@@ -302,4 +362,135 @@
                 raise ValueError("Requested %s nearest neighbors; acceptable numbers are integers greater than or equal to one, or None")
 
 
+    def __query_ball_point(self, x, r, p=2., eps=0):
+        R = Rectangle(self.maxes, self.mins)
 
+        def traverse_checking(node, rect):
+            if rect.min_distance_point(x,p)>=r/(1.+eps):
+                return []
+            elif rect.max_distance_point(x,p)<r*(1.+eps):
+                return traverse_no_checking(node)
+            elif isinstance(node, KDTree.leafnode):
+                d = self.data[node.idx]
+                return node.idx[distance(d,x,p)<=r].tolist()
+            else:
+                less, greater = rect.split(node.split_dim, node.split)
+                return traverse_checking(node.less, less)+traverse_checking(node.greater, greater)
+        def traverse_no_checking(node):
+            if isinstance(node, KDTree.leafnode):
+                
+                return node.idx.tolist()
+            else:
+                return traverse_no_checking(node.less)+traverse_no_checking(node.greater)
+
+        return traverse_checking(self.tree, R)
+
+    def query_ball_point(self, x, r, p=2., eps=0):
+        """Find all points within r of x
+
+        Parameters
+        ==========
+
+        x : array_like, shape tuple + (self.k,)
+            The point or points to search for neighbors of
+        r : positive float
+            The radius of points to return
+        p : float 1<=p<=infinity
+            Which Minkowski p-norm to use
+        eps : nonnegative float
+            Approximate search. Branches of the tree are not explored
+            if their nearest points are further than r/(1+eps), and branches
+            are added in bulk if their furthest points are nearer than r*(1+eps).
+
+        Returns
+        =======
+
+        results : list or array of lists
+            If x is a single point, returns a list of the indices of the neighbors
+            of x. If x is an array of points, returns an object array of shape tuple
+            containing lists of neighbors.
+
+
+        Note: if you have many points whose neighbors you want to find, you may save
+        substantial amounts of time by putting them in a KDTree and using query_ball_tree.
+        """
+        x = np.asarray(x)
+        if x.shape[-1]!=self.k:
+            raise ValueError("Searching for a %d-dimensional point in a %d-dimensional KDTree" % (x.shape[-1],self.k))
+        if len(x.shape)==1:
+            return self.__query_ball_point(x,r,p,eps)
+        else:
+            retshape = x.shape[:-1]
+            result = np.empty(retshape,dtype=np.object)
+            for c in np.ndindex(retshape):
+                result[c] = self.__query_ball_point(x[c], r, p=p, eps=eps)
+            return result
+
+    def query_ball_tree(self, other, r, p=2., eps=0):
+        """Find all pairs of points whose distance is at most r
+
+        Parameters
+        ==========
+
+        other : KDTree
+            The tree containing points to search against
+        r : positive float
+            The maximum distance
+        p : float 1<=p<=infinity
+            Which Minkowski norm to use
+        eps : nonnegative float
+            Approximate search. Branches of the tree are not explored
+            if their nearest points are further than r/(1+eps), and branches
+            are added in bulk if their furthest points are nearer than r*(1+eps).
+        
+        Returns
+        =======
+
+        results : list of lists
+            For each element self.data[i] of this tree, results[i] is a list of the
+            indices of its neighbors in other.data.
+        """
+        results = [[] for i in range(self.n)]
+        def traverse_checking(node1, rect1, node2, rect2):
+            if rect1.min_distance_rectangle(rect2, p)>r/(1.+eps):
+                return
+            elif rect1.max_distance_rectangle(rect2, p)<r*(1.+eps):
+                traverse_no_checking(node1, node2)
+            elif isinstance(node1, KDTree.leafnode):
+                if isinstance(node2, KDTree.leafnode):
+                    d = other.data[node2.idx]
+                    for i in node1.idx:
+                        results[i] += node2.idx[distance(d,self.data[i],p)<=r].tolist()
+                else:
+                    less, greater = rect2.split(node2.split_dim, node2.split)
+                    traverse_checking(node1,rect1,node2.less,less)
+                    traverse_checking(node1,rect1,node2.greater,greater)
+            elif isinstance(node2, KDTree.leafnode):
+                less, greater = rect1.split(node1.split_dim, node1.split)
+                traverse_checking(node1.less,less,node2,rect2)
+                traverse_checking(node1.greater,greater,node2,rect2)
+            else:
+                less1, greater1 = rect1.split(node1.split_dim, node1.split)
+                less2, greater2 = rect2.split(node2.split_dim, node2.split)
+                traverse_checking(node1.less,less1,node2.less,less2)
+                traverse_checking(node1.less,less1,node2.greater,greater2)
+                traverse_checking(node1.greater,greater1,node2.less,less2)
+                traverse_checking(node1.greater,greater1,node2.greater,greater2)
+
+        def traverse_no_checking(node1, node2):
+            if isinstance(node1, KDTree.leafnode):
+                if isinstance(node2, KDTree.leafnode):
+                    for i in node1.idx:
+                        results[i] += node2.idx.tolist()
+                else:
+                    traverse_no_checking(node1, node2.less)
+                    traverse_no_checking(node1, node2.greater)
+            else:
+                traverse_no_checking(node1.less, node2)
+                traverse_no_checking(node1.greater, node2)
+
+        traverse_checking(self.tree, Rectangle(self.maxes, self.mins),
+                          other.tree, Rectangle(other.maxes, other.mins))
+        return results
+
+        

Modified: branches/spatial/scipy/spatial/tests/test_kdtree.py
===================================================================
--- branches/spatial/scipy/spatial/tests/test_kdtree.py	2008-10-05 05:02:32 UTC (rev 4776)
+++ branches/spatial/scipy/spatial/tests/test_kdtree.py	2008-10-05 05:51:47 UTC (rev 4777)
@@ -3,7 +3,7 @@
 from numpy.testing import *
 
 import numpy as np
-from scipy.spatial import KDTree, distance
+from scipy.spatial import KDTree, distance, Rectangle
 
 class ConsistencyTests:
     def test_nearest(self):
@@ -126,7 +126,7 @@
         self.m = 4
 
 
-class CheckVectorization(NumpyTestCase):
+class test_vectorization:
     def setUp(self):
         self.data = np.array([[0,0,0],
                               [0,0,1],
@@ -176,4 +176,138 @@
         assert isinstance(d[0,0],list)
         assert isinstance(i[0,0],list)
 
+class ball_consistency:
 
+    def test_in_ball(self):
+        l = self.T.query_ball_point(self.x, self.d, p=self.p, eps=self.eps)
+        for i in l:
+            assert distance(self.data[i],self.x,self.p)<=self.d*(1.+self.eps)
+
+    def test_found_all(self):
+        c = np.ones(self.T.n,dtype=np.bool)
+        l = self.T.query_ball_point(self.x, self.d, p=self.p, eps=self.eps)
+        c[l] = False
+        assert np.all(distance(self.data[c],self.x,self.p)>=self.d/(1.+self.eps))
+
+class test_random_ball(ball_consistency):
+
+    def setUp(self):
+        n = 1000
+        k = 4
+        self.data = np.random.randn(n,k)
+        self.T = KDTree(self.data)
+        self.x = np.random.randn(k)
+        self.p = 2.
+        self.eps = 0
+        self.d = 0.2
+
+class test_random_ball_approx(test_random_ball):
+
+    def setUp(self):
+        test_random_ball.setUp(self)
+        self.eps = 0.1
+
+class test_random_ball_far(test_random_ball):
+
+    def setUp(self):
+        test_random_ball.setUp(self)
+        self.d = 2.
+
+class test_random_ball_l1(test_random_ball):
+
+    def setUp(self):
+        test_random_ball.setUp(self)
+        self.p = 1
+
+class test_random_ball_linf(test_random_ball):
+
+    def setUp(self):
+        test_random_ball.setUp(self)
+        self.p = np.inf
+
+def test_random_ball_vectorized():
+
+    n = 20
+    k = 5
+    T = KDTree(np.random.randn(n,k))
+    
+    r = T.query_ball_point(np.random.randn(2,3,k),1)
+    assert_equal(r.shape,(2,3))
+    assert isinstance(r[0,0],list)
+
+class two_trees_consistency:
+
+    def test_all_in_ball(self):
+        r = self.T1.query_ball_tree(self.T2, self.d, p=self.p, eps=self.eps)
+        for i, l in enumerate(r):
+            for j in l:
+                assert distance(self.data1[i],self.data2[j],self.p)<=self.d*(1.+self.eps)
+    def test_found_all(self):
+        r = self.T1.query_ball_tree(self.T2, self.d, p=self.p, eps=self.eps)
+        for i, l in enumerate(r):
+            c = np.ones(self.T2.n,dtype=np.bool)
+            c[l] = False
+            assert np.all(distance(self.data2[c],self.data1[i],self.p)>=self.d/(1.+self.eps))
+
+class test_two_random_trees(two_trees_consistency):
+
+    def setUp(self):
+        n = 100
+        k = 4
+        self.data1 = np.random.randn(n,k)
+        self.T1 = KDTree(self.data1,leafsize=2)
+        self.data2 = np.random.randn(n,k)
+        self.T2 = KDTree(self.data2,leafsize=2)
+        self.p = 2.
+        self.eps = 0
+        self.d = 0.2
+
+class test_two_random_trees_far(test_two_random_trees):
+
+    def setUp(self):
+        test_two_random_trees.setUp(self)
+        self.d = 2
+
+class test_two_random_trees_linf(test_two_random_trees):
+
+    def setUp(self):
+        test_two_random_trees.setUp(self)
+        self.p = np.inf
+
+
+class test_rectangle:
+
+    def setUp(self):
+        self.rect = Rectangle([0,0],[1,1])
+
+    def test_min_inside(self):
+        assert_almost_equal(self.rect.min_distance_point([0.5,0.5]),0)
+    def test_min_one_side(self):
+        assert_almost_equal(self.rect.min_distance_point([0.5,1.5]),0.5)
+    def test_min_two_sides(self):
+        assert_almost_equal(self.rect.min_distance_point([2,2]),np.sqrt(2))
+    def test_max_inside(self):
+        assert_almost_equal(self.rect.max_distance_point([0.5,0.5]),1/np.sqrt(2))
+    def test_max_one_side(self):
+        assert_almost_equal(self.rect.max_distance_point([0.5,1.5]),np.hypot(0.5,1.5))
+    def test_max_two_sides(self):
+        assert_almost_equal(self.rect.max_distance_point([2,2]),2*np.sqrt(2))
+
+    def test_split(self):
+        less, greater = self.rect.split(0,0.1)
+        assert_array_equal(less.maxes,[0.1,1])
+        assert_array_equal(less.mins,[0,0])
+        assert_array_equal(greater.maxes,[1,1])
+        assert_array_equal(greater.mins,[0.1,0])
+
+
+def test_distance_l2():
+    assert_almost_equal(distance([0,0],[1,1],2),np.sqrt(2))
+def test_distance_l1():
+    assert_almost_equal(distance([0,0],[1,1],1),2)
+def test_distance_linf():
+    assert_almost_equal(distance([0,0],[1,1],np.inf),1)
+def test_distance_vectorization():
+    x = np.random.randn(10,1,3)
+    y = np.random.randn(1,7,3)
+    assert_equal(distance(x,y).shape,(10,7))



More information about the Scipy-svn mailing list