How to delete elements from a numpy array using indecies returned by scipy.spatial.KDTree.query_ball_point method

Question

I am trying to use Kdtree data structure to remove closest points from an array preferablly without for loops.

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()  

but no matter what i do i can't get np.delete to work with the indecies returned by the ball_query method. What am i missing?

I would like to make the red colored points vanish in each iteration from the points array.

Answer 1

Your points array is a Nx2 matrix. Your ind indices are a list of row indices. What you need is to specify the axis along which you need deletion, ultimately this:

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

Also, once you delete indices, watch out for missing indices in your next iteration/calculations. Maybe you want to have a copy to delete points and plot and another copy for calculations that you do not delete from it.