simulation_object.py 文件源码

def calc_vbar(self):
        # Calculates the average velocities from neighboring birds depending on max_dist and max_num.
        my_tree = spatial.cKDTree(self.current_points)
        dist, indexes = my_tree.query(self.current_points, k=self._max_num)

        ri = np.zeros((len(self.current_points), self._max_num), dtype=int)
        rk = np.zeros_like(ri, dtype=int)

        good_inds = (dist < self._max_dist)
        ri[good_inds] = indexes[good_inds]
        rk[good_inds] = 1

        # I should get the angle and average
        ori = np.arctan2(self.current_v[:, 1], self.current_v[:, 0])

        mean_n = []
        for i in range(len(ri)):
            nei = ri[i][np.where(rk[i] == 1)[0]]
            mm = (np.arctan2(np.sum(np.sin(ori[nei])), np.sum(np.cos(ori[nei])))) % (2 * np.pi)
            mean_n.append(mm)

        vbar = np.array(zip(np.cos(mean_n), np.sin(mean_n)))
        return vbar, [np.array(ri), np.array(rk)]