Source Code for Module QSTK.qstklearn.kdtknn

  1  """ 
  2  A simple wrapper for scipy.spatial.kdtree.KDTree for doing KNN 
  3  """ 
  4  import math,random,sys,bisect,time 
  5  import numpy,scipy.spatial.distance 
  6  from scipy.spatial import cKDTree 
  7  import cProfile,pstats,gendata 
  8  import numpy as np 
  9   
 10 -class kdtknn(object): 
 11          """ 
 12          A simple wrapper of scipy.spatial.kdtree.KDTree 
 13           
 14          Since the scipy KDTree implementation does not allow for incrementally adding 
 15          data points, the entire KD-tree is rebuilt on the first call to 'query' after a 
 16          call to 'addEvidence'. For this reason it is more efficient to add training data 
 17          in batches. 
 18          """ 
 19 -        def __init__(self,k=3,method='mean',leafsize=10): 
 20                  """ 
 21                  Basic setup. 
 22                  """ 
 23                  self.leafsize = leafsize 
 24                  self.data = None 
 25                  self.kdt = None 
 26                  self.rebuild_tree = True 
 27                  self.k = k 
 28                  self.method = method 
 29   
 30 -        def addEvidence(self,dataX,dataY=None): 
 31                  """ 
 32                  @summary: Add training data 
 33                  @param dataX: Data to add, either entire set with classification as last column, or not if 
 34                                the Y data is provided explicitly.  Must be same width as previously appended data. 
 35                  @param dataY: Optional, can be used  
 36                   
 37                  'data' should be a numpy array matching the same dimensions as any data  
 38                  provided in previous calls to addEvidence, with dataY as the  
 39                  training label. 
 40                  """ 
 41                   
 42                  ''' Slap on Y column if it is provided, if not assume it is there ''' 
 43                  if not dataY == None: 
 44                          data = numpy.zeros([dataX.shape[0],dataX.shape[1]+1]) 
 45                          data[:,0:dataX.shape[1]]=dataX 
 46                          data[:,(dataX.shape[1])]=dataY 
 47                  else: 
 48                          data = dataX 
 49                   
 50                  self.rebuild_tree = True 
 51                  if self.data is None: 
 52                          self.data = data 
 53                  else: 
 54                          self.data = numpy.append(self.data,data,axis=0) 
 55   
 56 -        def rebuildKDT(self): 
 57                  """ 
 58                  Force the internal KDTree to be rebuilt. 
 59                  """ 
 60                  self.kdt = cKDTree(self.data[:,:-1],leafsize=self.leafsize) 
 61                  self.rebuild_tree = False 
 62           
 63 -        def query(self,points,k=None,method=None): 
 64                  """ 
 65                  Classify a set of test points given their k nearest neighbors. 
 66                   
 67                  'points' should be a numpy array with each row corresponding to a specific query. 
 68                  Returns the estimated class according to supplied method (currently only 'mode' 
 69                  and 'mean' are supported) 
 70                  """ 
 71                  if k is None: 
 72                          k = self.k 
 73                  if method is None: 
 74                          method = self.method 
 75                  if self.rebuild_tree is True: 
 76                          if self.data is None: 
 77                                  return None 
 78                          self.rebuildKDT() 
 79                  #kdt.query returns a list of distances and a list of indexes into the 
 80                  #data array 
 81                  tmp = self.kdt.query(points,k) 
 82                  if k == 1: 
 83                          #in the case of k==1, numpy fudges an array of 1 dimension into 
 84                          #a scalar, so we handle it seperately. tmp[1] is the list of 
 85                          #indecies, tmp[1][0] is the first one (we only need one), 
 86                          #self.data[tmp[1][0]] is the data point corresponding to the 
 87                          #first neighbor, and self.data[tmp[1][0]][-1] is the last column 
 88                          #which is the class of the neighbor. 
 89                          return self.data[tmp[1][0]][-1] 
 90                  #for all the neighbors returned by kdt.query, get their class and stick that into a list 
 91                  n_clsses = map(lambda rslt: map(lambda p: p[-1], self.data[rslt]), self.kdt.query(points,k)[1]) 
 92                  if method=='mode': 
 93                          return map(lambda x: scipy.stats.stats.mode(x)[0],n_clsses)[0] 
 94                  elif method=='mean': 
 95                          return numpy.array(map(lambda x: numpy.mean(x),n_clsses)) 
 96                  elif method=='median': 
 97                          return numpy.array(map(lambda x: numpy.median(x),n_clsses)) 
 98                  elif method=='raw': 
 99                          return numpy.array(n_clsses) 
100   
101 -def getflatcsv(fname): 
102          inf = open(fname) 
103          return numpy.array([map(float,s.strip().split(',')) for s in inf.readlines()]) 
104   
105 -def testgendata(): 
106          fname = 'test2.dat' 
107          querys = 1000 
108          d = 2 
109          k=3 
110          bnds = ((-10,10),)*d 
111          clsses = (0,1) 
112          data = getflatcsv(fname) 
113          kdt = kdtknn(k,method='mode') 
114          kdt.addEvidence(data) 
115          kdt.rebuildKDT() 
116          stime = time.time() 
117          for x in xrange(querys): 
118                  pnt = numpy.array(gendata.gensingle(d,bnds,clsses)) 
119                  reslt = kdt.query(numpy.array([pnt[:-1]])) 
120                  print pnt,"->",reslt 
121          etime = time.time() 
122          print etime-stime,'/',querys,'=',(etime-stime)/float(querys),'avg wallclock time per query' 
123          #foo.addEvidence(data[:,:-1],data[:,-1]) 
124          #foo.num_checks = 0 
125          #for x in xrange(querys): 
126          #       pnt = numpy.array(gendata.gensingle(d,bnds,clsses)) 
127          #       foo.query(pnt[:-1],3) 
128          #       if x % 50 == 0: 
129          #               print float(foo.num_checks)/float(x+1), 
130          #               print x,"/",querys 
131          #print "Average # queries:", float(foo.num_checks)/float(querys) 
132           
133 -def test(): 
134          testgendata() 
135   
136  if __name__=="__main__": 
137          test() 
138          #prof= cProfile.Profile() 
139          #prof.run('test()') 
140          #stats = pstats.Stats(prof) 
141          #stats.sort_stats("cumulative").print_stats() 
142