def _setup_potential(self, R, z, use_pkdgrav = False) :
from galpy.potential import vcirc
# cast the points into arrays for compatibility
if isinstance(R,float) :
R = np.array([R])
if isinstance(z, float) :
z = np.array([z])
# compute the hash for the requested grid
new_hash = hashlib.md5(np.array([R,z])).hexdigest()
# if we computed for these points before, return; otherwise compute
if new_hash in self._point_hash :
pot, r_acc = self._point_hash[new_hash]
# if use_pkdgrav :
else :
# set up the four points per R,z pair to mimic axisymmetry
points = np.zeros((len(R),len(z),4,3))
for i in xrange(len(R)) :
for j in xrange(len(z)) :
points[i,j] = [(R[i],0,z[j]),
(0,R[i],z[j]),
(-R[i],0,z[j]),
(0,-R[i],z[j])]
points_new = points.reshape(points.size/3,3)
pot, acc = grav_omp.direct(self._s,points_new,num_threads=self._num_threads)
pot = pot.reshape(len(R)*len(z),4)
acc = acc.reshape(len(R)*len(z),4,3)
# need to average the potentials
if len(pot) > 1:
pot = pot.mean(axis=1)
else :
pot = pot.mean()
# get the radial accelerations
r_acc = np.zeros((len(R)*len(z),2))
rvecs = [(1.0,0.0,0.0),
(0.0,1.0,0.0),
(-1.0,0.0,0.0),
(0.0,-1.0,0.0)]
# reshape the acc to make sure we have a leading index even
# if we are only evaluating a single point, i.e. we have
# shape = (1,4,3) not (4,3)
acc = acc.reshape((len(r_acc),4,3))
for i in xrange(len(R)) :
for j,rvec in enumerate(rvecs) :
r_acc[i,0] += acc[i,j].dot(rvec)
r_acc[i,1] += acc[i,j,2]
r_acc /= 4.0
# store the computed values for reuse
self._point_hash[new_hash] = [pot,r_acc]
return pot, r_acc
def _setup_potential(self, R, z, use_pkdgrav = False, dr = 0.01) :
"""
Calculates the potential and force grids for the snapshot for
use with other galpy functions.
**Input**:
*R*: R grid coordinates
*z*: z grid coordinates
**Optional Keywords**:
*use_pkdgrav*: (False) whether to use pkdgrav for the gravity
calculation
*dr*: (0.01) offset to use for the gradient calculation - the
points are positioned at +/- dr from the central point
"""
from galpy.potential import vcirc
# cast the points into arrays for compatibility
if isinstance(R,float) :
R = np.array([R])
if isinstance(z, float) :
z = np.array([z])
# set up the four points per R,z pair to mimic axisymmetry
points = np.zeros((len(R),len(z),4,3))
for i in xrange(len(R)) :
for j in xrange(len(z)) :
points[i,j] = [(R[i],0,z[j]),
(0,R[i],z[j]),
(-R[i],0,z[j]),
(0,-R[i],z[j])]
points_new = points.reshape(points.size/3,3)
self.points = points_new
# set up the points to calculate the second derivatives
zgrad_points = np.zeros((len(points_new)*2,3))
rgrad_points = np.zeros((len(points_new)*2,3))
for i,p in enumerate(points_new) :
zgrad_points[i*2] = p
zgrad_points[i*2][2] -= dr
zgrad_points[i*2+1] = p
zgrad_points[i*2+1][2] += dr
rgrad_points[i*2] = p
rgrad_points[i*2][:2] -= p[:2]/np.sqrt(np.dot(p[:2],p[:2]))*dr
rgrad_points[i*2+1] = p
rgrad_points[i*2+1][:2] += p[:2]/np.sqrt(np.dot(p[:2],p[:2]))*dr
if use_pkdgrav :
raise RuntimeError("using pkdgrav not currently implemented")
sn = pynbody.snapshot._new(len(self._s.d)+len(self._s.g)+len(self._s.s)+len(points_new))
print "setting up %d grid points"%(len(points_new))
#sn['pos'][0:len(self.s)] = self.s['pos']
#sn['mass'][0:len(self.s)] = self.s['mass']
#sn['phi'] = 0.0
#sn['eps'] = 1e3
#sn['eps'][0:len(self.s)] = self.s['eps']
#sn['vel'][0:len(self.s)] = self.s['vel']
#sn['mass'][len(self.s):] = 1e-10
sn['pos'][len(self._s):] = points_new
sn['mass'][len(self._s):] = 0.0
sn.write(fmt=pynbody.tipsy.TipsySnap, filename='potgridsnap')
command = '~/bin/pkdgrav2_pthread -sz %d -n 0 +std -o potgridsnap -I potgridsnap +potout +overwrite %s'%(self._numcores, self._s._paramfile['filename'])
print command
system(command)
sn = pynbody.load('potgridsnap')
acc = sn['accg'][len(self._s):].reshape(len(R)*len(z),4,3)
pot = sn['pot'][len(self._s):].reshape(len(R)*len(z),4)
else :
if self._interpPot:
pot, acc = grav_omp.direct(self._s,points_new,num_threads=self._numcores)
pot = pot.reshape(len(R)*len(z),4)
acc = acc.reshape(len(R)*len(z),4,3)
# need to average the potentials
if len(pot) > 1:
pot = pot.mean(axis=1)
else :
pot = pot.mean()
# get the radial accelerations
rz_acc = np.zeros((len(R)*len(z),2))
rvecs = [(1.0,0.0,0.0),
(0.0,1.0,0.0),
(-1.0,0.0,0.0),
(0.0,-1.0,0.0)]
# reshape the acc to make sure we have a leading index even
# if we are only evaluating a single point, i.e. we have
# shape = (1,4,3) not (4,3)
acc = acc.reshape((len(rz_acc),4,3))
for i in xrange(len(R)*len(z)) :
for j,rvec in enumerate(rvecs) :
rz_acc[i,0] += acc[i,j].dot(rvec)
rz_acc[i,1] += acc[i,j,2]
rz_acc /= 4.0
self._potGrid = pot.reshape((len(R),len(z)))
self._rforceGrid = rz_acc[:,0].reshape((len(R),len(z)))
self._zforceGrid = rz_acc[:,1].reshape((len(R),len(z)))
# compute the force gradients
# first get the accelerations
if self._interpverticalfreq :
zgrad_pot, zgrad_acc = grav_omp.direct(self._s,zgrad_points,num_threads=self._numcores)
# each point from the points used above for pot and acc is straddled by
# two points to get the gradient across it. Compute the gradient by
# using a finite difference
zgrad = np.zeros(len(points_new))
# do a loop through the pairs of points -- reshape the array
# so that each item is the pair of acceleration vectors
# then calculate the gradient from the two points
for i,zacc in enumerate(zgrad_acc.reshape((len(zgrad_acc)/2,2,3))) :
zgrad[i] = ((zacc[1]-zacc[0])/(dr*2.0))[2]
# reshape the arrays
self._z2derivGrid = zgrad.reshape((len(zgrad)/4,4)).mean(axis=1).reshape((len(R),len(z)))
# do the same for the radial component
if self._interpepifreq:
rgrad_pot, rgrad_acc = grav_omp.direct(self._s,rgrad_points,num_threads=self._numcores)
rgrad = np.zeros(len(points_new))
for i,racc in enumerate(rgrad_acc.reshape((len(rgrad_acc)/2,2,3))) :
point = points_new[i]
point[2] = 0.0
rvec = point/np.sqrt(np.dot(point,point))
rgrad_vec = (np.dot(racc[1],rvec)-
np.dot(racc[0],rvec)) / (dr*2.0)
rgrad[i] = rgrad_vec
self._R2derivGrid = rgrad.reshape((len(rgrad)/4,4)).mean(axis=1).reshape((len(R),len(z)))
