two_point_functions.py 文件源码

def _find_nearest_cell(self, points):
        """
        Finds the closest grid cell for each point in a vectorized manner.
        """
        if self.nlevels == 0:
            pos = (points - self.ds.left_edge) / self.width
            n = (self.sizes[2] * pos[:,2]).astype('int32')
            n += self.sizes[2] * (self.sizes[1] * pos[:,1]).astype('int32')
            n += self.sizes[2] * self.sizes[1] * (self.sizes[0] * pos[:,0]).astype('int32')
        else:
            # Normalize the points to a 1-period for use only within the kdtree.
            points[:, 0] = points[:, 0] / self.period[0]
            points[:, 1] = points[:, 1] / self.period[1]
            points[:, 2] = points[:, 2] / self.period[2]
            fKD.qv_many = points.T
            fKD.nn_tags = np.asfortranarray(np.empty((1, points.shape[0]), dtype='int64'))
            fKD.find_many_nn_nearest_neighbors()
            # The -1 is for fortran counting.
            n = fKD.nn_tags[0,:] - 1
        return n