from numarray import *
import SM

from pNbody import *

import time

import treelib


eps = 0.1

DesNumNgb = 32
MaxNumNgbDeviation = 2

# open a model
nb = Nbody('disk.dat',ftype='gadget')
nb = nb.select('wg')


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


R = nb.rxy()
pos = nb.pos


# create a tree object
print "build tree"
MyTree = nb.getTree()

# init hsml
print "init hsml"
#nb.hsml = MyTree.InitHsml(DesNumNgb,MaxNumNgbDeviation)
#nb.rsp = nb.get_rsp_approximation(DesNumNgb,MaxNumNgbDeviation)

# evaluate density (all particles)
#print "compute density for all particles"
#nb.density,nb.hsml = MyTree.Density(nb.pos,nb.hsml,DesNumNgb,MaxNumNgbDeviation)
nb.rho = nb.rho
nb.rsp = nb.rsp

# evaluate density
print "compute density and hsml for points"
hsml = ones(len(pos)).astype(Float32)
density,hsml = MyTree.Density(pos,hsml,DesNumNgb,MaxNumNgbDeviation)
#density = nb.density
#hsml = nb.hsml

# evaluate
print "evaluate observable"
nb.val = nb.Vr()
val1 = MyTree.SphEvaluate(pos,hsml,nb.density,nb.val,DesNumNgb,MaxNumNgbDeviation)

nb.val = nb.Vr()**2
val2 = MyTree.SphEvaluate(pos,hsml,nb.density,nb.val,DesNumNgb,MaxNumNgbDeviation)

val = sqrt(val2-val1**2)


#nb.opt = val
#nb.write()



g = SM.plot()
g.limits(R,0,0.5)
g.box()
g.points(R,val)
g.show()

