如何使用scipy.spatial.KDTree.query_ball_point方法返回的索引从numpy数组中删除元素

Question

我正在尝试使用Kdtree数据结构来从数组中移除最接近的点，而不需要for循环。

import sys

import time

import scipy.spatial

class KDTree:
    """
    Nearest neighbor search class with KDTree
    """

    def __init__(self, data):
        # store kd-tree
        self.tree = scipy.spatial.cKDTree(data)

    def search(self, inp, k=1):
        """
        Search NN
        inp: input data, single frame or multi frame
        """

        if len(inp.shape) >= 2:  # multi input
            index = []
            dist = []

            for i in inp.T:
                idist, iindex = self.tree.query(i, k=k)
                index.append(iindex)
                dist.append(idist)

            return index, dist

        dist, index = self.tree.query(inp, k=k)
        return index, dist

    def search_in_distance(self, inp, r):
        """
        find points with in a distance r
        """

        index = self.tree.query_ball_point(inp, r)
        return np.asarray(index)


import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
start = time.time()
fig, ar = plt.subplots()
t = 0
R = 50.0
u = R *np.cos(t)
v = R *np.sin(t)

x = np.linspace(-100,100,51)
y = np.linspace(-100,100,51)

xx, yy = np.meshgrid(x,y)
points =np.vstack((xx.ravel(),yy.ravel())).T
Tree = KDTree(points)
ind = Tree.search_in_distance([u, v],10.0)
ar.scatter(points[:,0],points[:,1],c='k',s=1)
infected = points[ind]
ar.scatter(infected[:,0],infected[:,1],c='r',s=5)

def animate(i):
    global R,t,start,points
    ar.clear()
    u = R *np.cos(t)
    v = R *np.sin(t)
    ind = Tree.search_in_distance([u, v],10.0)
    ar.scatter(points[:,0],points[:,1],c='k',s=1)
    infected = points[ind]
    ar.scatter(infected[:,0],infected[:,1],c='r',s=5)
    #points = np.delete(points,ind)
    t+=0.01
    end = time.time()
    if end - start != 0:
        print((end - start), end="\r")
        start = end
ani = animation.FuncAnimation(fig, animate, interval=20)
plt.show()  

但是无论我做什么，我都无法np.delete来处理ball_query方法返回的索引。我想念什么？

我想使红色点在每次迭代中都从点数组中消失。

Answer 1

您的points数组是Nx2矩阵。您的索引是行索引的列表。您需要指定要删除的轴，最终是这样：

points = np.delete(points,ind,axis=0)

此外，一旦删除索引，请在下次迭代/计算中注意缺少的索引。也许您想要一个副本来删除点和图，并想要另一个副本来进行计算，但不要从中删除。