Ejemplo n.º 1
0
    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
Ejemplo n.º 2
0
    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)))