In [23]:

    

from scipy.spatial import cKDTree, KDTree
import numpy as np
from time import clock

np.random.seed(1337)

def test():
    x, y = np.mgrid[0:10, 0:10]
    points = np.array(zip(x.ravel(), y.ravel()))
    tree = cKDTree(points)
    # print(tree.data)
    erg = tree.query_ball_point([3, 3], 1)
    # print(erg)
    # print(points[erg])

if __name__ == '__main__':
    t0 = clock()
    test()
    t1 = clock()
    print "Elapsed time %.2f seconds" % (t1-t0)


    

    




Elapsed time 0.00 seconds

In [27]:

    

particle_positions = np.random.random((int(1e6),2))
particle_positions_tree = cKDTree(particle_positions)

r = 1e-2
print((1/r)**2)
# print(np.sqrt(1e6))
%timeit erg = np.array(particle_positions_tree.query_ball_point([0.5, 0.5], r))
erg = np.array(particle_positions_tree.query_ball_point([0.5, 0.5], r))

print(len(erg))
print(particle_positions[erg])


    

    




10000.0
10000 loops, best of 3: 49.5 µs per loop
317
[[ 0.49423924  0.49198428]
 [ 0.49512143  0.49180124]
 [ 0.49587472  0.49103919]
 [ 0.49529721  0.49136098]
 [ 0.4959342   0.49203431]
 [ 0.49648337  0.49116434]
 [ 0.4974648   0.49094373]
 [ 0.49713409  0.490448  ]
 [ 0.49935016  0.49155692]
 [ 0.49957841  0.4905388 ]
 [ 0.50092912  0.49042392]
 [ 0.49141325  0.4959088 ]
 [ 0.49178185  0.49489868]
 [ 0.49259434  0.49738631]
 [ 0.49296396  0.49333404]
 [ 0.49303758  0.49751649]
 [ 0.49151745  0.49777586]
 [ 0.49102412  0.49797701]
 [ 0.49160375  0.49813214]
 [ 0.49172292  0.49852746]
 [ 0.4916194   0.49825751]
 [ 0.49192123  0.49831187]
 [ 0.49091029  0.4985594 ]
 [ 0.49068859  0.49875248]
 [ 0.49151034  0.49905053]
 [ 0.4918668   0.49990346]
 [ 0.49207551  0.49959077]
 [ 0.49193354  0.49966143]
 [ 0.49156593  0.49878066]
 [ 0.49303245  0.49900163]
 [ 0.49190981  0.49907085]
 [ 0.4939823   0.49520416]
 [ 0.49420154  0.49523955]
 [ 0.49434923  0.49493262]
 [ 0.49454914  0.49369957]
 [ 0.49468549  0.49516258]
 [ 0.49498931  0.49522629]
 [ 0.49340706  0.49515122]
 [ 0.4939457   0.49532715]
 [ 0.49570838  0.49580843]
 [ 0.49461664  0.49539619]
 [ 0.49653535  0.4968458 ]
 [ 0.4974036   0.49344027]
 [ 0.49710516  0.4926906 ]
 [ 0.49660658  0.49609822]
 [ 0.49746646  0.49358099]
 [ 0.49756838  0.49486867]
 [ 0.49783916  0.49221258]
 [ 0.49848705  0.49428276]
 [ 0.49987063  0.49316932]
 [ 0.49976353  0.49468732]
 [ 0.49889039  0.4939798 ]
 [ 0.49985051  0.49246343]
 [ 0.49914966  0.49533909]
 [ 0.49915924  0.49618618]
 [ 0.49875593  0.49479102]
 [ 0.50008257  0.49669825]
 [ 0.49918856  0.49220049]
 [ 0.49924607  0.49533478]
 [ 0.49999354  0.49564249]
 [ 0.4997148   0.49574996]
 [ 0.49997578  0.4935545 ]
 [ 0.49469875  0.49686533]
 [ 0.49401274  0.4978833 ]
 [ 0.49411264  0.49793917]
 [ 0.49517672  0.497689  ]
 [ 0.49465254  0.49788986]
 [ 0.49516432  0.49798789]
 [ 0.49334559  0.49843407]
 [ 0.49510616  0.49939898]
 [ 0.49521113  0.49973164]
 [ 0.49529809  0.49893981]
 [ 0.49539752  0.49726284]
 [ 0.49556007  0.49799067]
 [ 0.49581629  0.4974524 ]
 [ 0.49606955  0.49964664]
 [ 0.49672686  0.49915529]
 [ 0.49699676  0.4995754 ]
 [ 0.49730001  0.49773803]
 [ 0.49980761  0.4975238 ]
 [ 0.49921623  0.49934192]
 [ 0.49835418  0.49883058]
 [ 0.49802624  0.49807664]
 [ 0.49740168  0.49919505]
 [ 0.49968339  0.49749647]
 [ 0.49795268  0.49936703]
 [ 0.50019171  0.49874689]
 [ 0.49894691  0.49908428]
 [ 0.50015537  0.49905977]
 [ 0.49974667  0.49696863]
 [ 0.50043261  0.49876142]
 [ 0.49872067  0.49766481]
 [ 0.49744176  0.49926815]
 [ 0.50099862  0.49074473]
 [ 0.50122037  0.49083378]
 [ 0.50240995  0.49144681]
 [ 0.50174522  0.49120725]
 [ 0.50153395  0.49172145]
 [ 0.50257712  0.49179319]
 [ 0.50184016  0.49248232]
 [ 0.50217887  0.49313975]
 [ 0.50214967  0.49340164]
 [ 0.50264726  0.49369859]
 [ 0.50299152  0.49089289]
 [ 0.50296267  0.49293699]
 [ 0.50293639  0.49331416]
 [ 0.50305829  0.49048811]
 [ 0.50346128  0.49266535]
 [ 0.50400616  0.49258254]
 [ 0.5043079   0.49199859]
 [ 0.50482424  0.49229931]
 [ 0.50550362  0.49181087]
 [ 0.50544255  0.49186895]
 [ 0.50499989  0.49142621]
 [ 0.5016941   0.49372467]
 [ 0.50160447  0.49405329]
 [ 0.5011788   0.49465951]
 [ 0.50301439  0.49462346]
 [ 0.50389866  0.49395784]
 [ 0.5041343   0.49436074]
 [ 0.50548708  0.49426343]
 [ 0.50430407  0.49424491]
 [ 0.50544032  0.49394729]
 [ 0.50552541  0.49520328]
 [ 0.50220004  0.49554964]
 [ 0.50213682  0.49537006]
 [ 0.50222104  0.49582529]
 [ 0.50173968  0.49598665]
 [ 0.50130408  0.49639348]
 [ 0.50518606  0.49660706]
 [ 0.50527904  0.49663067]
 [ 0.5051284   0.49670805]
 [ 0.50169246  0.49786145]
 [ 0.50303352  0.49682427]
 [ 0.50400383  0.49693859]
 [ 0.50208099  0.49796668]
 [ 0.5039797   0.49970697]
 [ 0.5015185   0.49971061]
 [ 0.50445853  0.49976255]
 [ 0.50180835  0.49964439]
 [ 0.50421957  0.49934798]
 [ 0.50212844  0.49920242]
 [ 0.50523136  0.49848703]
 [ 0.50188955  0.49849994]
 [ 0.50593454  0.49486207]
 [ 0.50601278  0.49498232]
 [ 0.50624597  0.4940468 ]
 [ 0.50683239  0.49410274]
 [ 0.50818407  0.49522245]
 [ 0.50823963  0.49577407]
 [ 0.50594226  0.4961132 ]
 [ 0.50776936  0.49531146]
 [ 0.50641226  0.4961336 ]
 [ 0.50839923  0.49658234]
 [ 0.50784059  0.49670042]
 [ 0.50748432  0.49698303]
 [ 0.50554583  0.49704269]
 [ 0.50916625  0.49710385]
 [ 0.50724319  0.49733988]
 [ 0.5074655   0.49793951]
 [ 0.50678589  0.49714073]
 [ 0.50829482  0.49806726]
 [ 0.5070675   0.49953369]
 [ 0.50860009  0.49849297]
 [ 0.5078748   0.49810754]
 [ 0.50583576  0.4998606 ]
 [ 0.50708065  0.49878318]
 [ 0.50693147  0.49980044]
 [ 0.50917028  0.49970715]
 [ 0.50729989  0.49818782]
 [ 0.50976298  0.49946622]
 [ 0.49072397  0.50227625]
 [ 0.49075992  0.50293677]
 [ 0.49087834  0.50090541]
 [ 0.49102696  0.50379731]
 [ 0.49109209  0.50378608]
 [ 0.49139316  0.50397168]
 [ 0.4920701   0.50308623]
 [ 0.4924563   0.50319603]
 [ 0.49235897  0.50041332]
 [ 0.49207231  0.50075126]
 [ 0.49189556  0.50014188]
 [ 0.49281383  0.50112268]
 [ 0.49147573  0.50351664]
 [ 0.4937317   0.50386739]
 [ 0.49336932  0.50440495]
 [ 0.49386272  0.50073996]
 [ 0.49399417  0.50093932]
 [ 0.49443882  0.50395908]
 [ 0.49485199  0.50024632]
 [ 0.49518399  0.5016584 ]
 [ 0.49555118  0.50131485]
 [ 0.49634744  0.50368582]
 [ 0.49648571  0.50451227]
 [ 0.49665891  0.50424808]
 [ 0.49682668  0.50225731]
 [ 0.4970811   0.50482665]
 [ 0.49729459  0.50301882]
 [ 0.4975545   0.50460356]
 [ 0.49764659  0.50004778]
 [ 0.49776454  0.50312937]
 [ 0.49803355  0.50405079]
 [ 0.49767037  0.50133742]
 [ 0.49263557  0.50635446]
 [ 0.49256388  0.50503647]
 [ 0.49197168  0.50489308]
 [ 0.49328964  0.50661797]
 [ 0.4940609   0.50604297]
 [ 0.49452595  0.50598605]
 [ 0.4944541   0.50649243]
 [ 0.4947452   0.50539722]
 [ 0.4951435   0.50582856]
 [ 0.49812914  0.50734116]
 [ 0.4956492   0.506987  ]
 [ 0.49647217  0.50545299]
 [ 0.49683766  0.50614414]
 [ 0.49780173  0.50693391]
 [ 0.49628217  0.50783668]
 [ 0.49668075  0.5076632 ]
 [ 0.49613103  0.50503375]
 [ 0.49720345  0.50819128]
 [ 0.49545837  0.50828087]
 [ 0.49522954  0.50591311]
 [ 0.49752686  0.50916953]
 [ 0.49704394  0.50894636]
 [ 0.49775364  0.50918963]
 [ 0.49787425  0.50876439]
 [ 0.49753185  0.5092587 ]
 [ 0.50072089  0.50255832]
 [ 0.49988954  0.50291632]
 [ 0.50077648  0.50300476]
 [ 0.50076867  0.50056341]
 [ 0.50073747  0.50028948]
 [ 0.50282743  0.5026837 ]
 [ 0.50255636  0.50371566]
 [ 0.50276846  0.50292582]
 [ 0.50287868  0.50232494]
 [ 0.5030363   0.5009073 ]
 [ 0.50136994  0.50354111]
 [ 0.50209703  0.5023962 ]
 [ 0.50135476  0.50243846]
 [ 0.50161083  0.50240958]
 [ 0.50365227  0.50116302]
 [ 0.50388191  0.50153795]
 [ 0.50407256  0.50238988]
 [ 0.50416981  0.5021385 ]
 [ 0.50412299  0.50199203]
 [ 0.50420327  0.50073701]
 [ 0.50619272  0.50194628]
 [ 0.50429818  0.50134021]
 [ 0.50512894  0.50144009]
 [ 0.50529489  0.50101785]
 [ 0.50571975  0.502566  ]
 [ 0.5052551   0.50163633]
 [ 0.50595667  0.50225101]
 [ 0.50645055  0.50028172]
 [ 0.50646308  0.50249045]
 [ 0.50653945  0.50390398]
 [ 0.50205359  0.50409584]
 [ 0.50060334  0.50468255]
 [ 0.49899311  0.50469601]
 [ 0.50106691  0.50505058]
 [ 0.499397    0.50779516]
 [ 0.49922     0.50766983]
 [ 0.50154468  0.50579604]
 [ 0.50174722  0.5068226 ]
 [ 0.5015315   0.50713862]
 [ 0.50163838  0.50671928]
 [ 0.49958498  0.50653653]
 [ 0.50002805  0.50584809]
 [ 0.50009235  0.50442654]
 [ 0.50213169  0.50435144]
 [ 0.50219256  0.50718303]
 [ 0.50240515  0.50443569]
 [ 0.50263114  0.50757908]
 [ 0.50357866  0.50578449]
 [ 0.50373726  0.50565375]
 [ 0.50422411  0.50538507]
 [ 0.50546237  0.50710529]
 [ 0.5045181   0.50537039]
 [ 0.50598373  0.50790644]
 [ 0.50623039  0.50763265]
 [ 0.50635694  0.50539146]
 [ 0.50638436  0.50764697]
 [ 0.50518168  0.50493244]
 [ 0.50483896  0.5054513 ]
 [ 0.50646403  0.50701404]
 [ 0.50565123  0.50791873]
 [ 0.50660758  0.50336643]
 [ 0.5068884   0.50036361]
 [ 0.50706481  0.50331802]
 [ 0.50717547  0.50008066]
 [ 0.50753531  0.50105928]
 [ 0.50781074  0.50310836]
 [ 0.50802821  0.50050527]
 [ 0.50891045  0.4999744 ]
 [ 0.50814197  0.50089728]
 [ 0.50896357  0.50014801]
 [ 0.50906642  0.50017284]
 [ 0.50913301  0.50322759]
 [ 0.50978096  0.5013268 ]
 [ 0.50948994  0.50134189]
 [ 0.50722677  0.5050204 ]
 [ 0.50716008  0.50430578]
 [ 0.50708132  0.50383006]
 [ 0.50735975  0.50577522]
 [ 0.49978351  0.50836168]
 [ 0.49900892  0.50968975]
 [ 0.49994012  0.50828169]
 [ 0.50014645  0.50988733]
 [ 0.50131     0.50835887]
 [ 0.50048744  0.50934633]
 [ 0.50224075  0.50899933]
 [ 0.50254375  0.50844087]
 [ 0.5025518   0.50814493]
 [ 0.50427896  0.50895296]
 [ 0.50335174  0.50902822]]

In [3]:

    

exit(1)


    

In [ ]: