from numarray import *
import SM
import treelib

from pNbody import *

import time



eps = 0.1

# open a model
nb = Nbody('disk.dat',ftype='gadget')
#nb = nb.selectc(nb.rxyz()<2)

# create a tree object
MyTree = treelib.Tree(npart=nb.npart,pos=nb.pos,mass=nb.mass)






R = arange(0,20,0.1)
print len(R)
pos = concatenate((R,zeros(len(R)),zeros(len(R))))
pos.shape = (3,len(R))
pos = transpose(pos)
pos = pos.astype(Float32)



t1 = time.time()
pot = MyTree.Potential(pos,eps)
t2 = time.time()
pot = pot.astype(Float32)

print t2-t1

pots2 = array([],Float)
t3 = time.time()
for ppos in pos:
  #pots1 = concatenate((pots1, MyTree.Potential([ppos],eps)))  
  pots2 = concatenate((pots2,nb.Pot(ppos,eps)))  
t4 = time.time()



print t4-t3


R = R.astype(Float32)

g = SM.plot()
#g.limits(R,pots2)
g.limits(R,pot)
g.box()
g.connect(R,pots2)
#g.connect(R,pots1)
g.connect(R,pot)
g.show()

