Removig points that are close to each other

remove_close_points.py

"""
Прореживание облака 2D-точек по минимальному расстоянию между точками

Реализация "в лоб" с использованием KD-Tree
"""

import numpy as np
import scipy.spatial as spatial


def remove_close_points(px, py, mindist):
    points = np.array(list(zip(px, py)))
    kdtree = spatial.cKDTree(points)

    exclude = {}

    for i, pt in enumerate(points):
        if i in exclude:
            continue

        nhoods = kdtree.query_ball_point(pt, mindist)
        exclude.update({n: None for n in nhoods if n != i})

    exclude = list(exclude.keys())
    res_points = np.delete(points, exclude, axis=0)

    x = res_points[:, 0]
    y = res_points[:, 1]

    return x, y

remove_close_points2.py

"""
Прореживание облака 2D-точек по минимальному расстоянию между точками

Реализация на основе равномерной сетки с использованием KD-Tree
При большом количестве точек работает быстрее чем реализация "в лоб", но имеет 
свойство генерировать не очень "красивое" распределение точек в прореженном облаке. 
Точки стремятся выстроиться по сетке.
"""

import itertools

import numpy as np
import scipy.spatial as spatial

def remove_close_points2(px, py, mindist):

    # Формируем равномерную сетку узлов с шагом не менее mindist
    xmin, xmax = np.min(px), np.max(px)
    ymin, ymax = np.min(py), np.max(py)

    gx_count = np.ceil(np.abs(xmax - xmin) / mindist)
    gy_count = np.ceil(np.abs(ymax - ymin) / mindist)

    gx = np.linspace(xmin, xmax, gx_count)
    gy = np.linspace(ymin, ymax, gy_count)
    grid = list(itertools.product(gx, gy))

    points = np.array(list(zip(px, py)))
    kdtree = spatial.cKDTree(points)

    # Для каждой точки сетки с помощью KD-Tree ищем ближайшую к ней точку из
    # всего множества точек на расстоянии не более mindist
    gd, idx = kdtree.query(grid, k=1, distance_upper_bound=mindist, n_jobs=-1)

    idx = idx[gd != np.inf]
    rare_points = points[idx]

    kdtree = spatial.cKDTree(rare_points)
    exclude = {}

    for i, pt in enumerate(rare_points):
        if i in exclude:
            continue

        nhoods = kdtree.query_ball_point(pt, mindist)
        exclude.update({n: None for n in nhoods if n != i})

    exclude = list(exclude.keys())
    res_points = np.delete(rare_points, exclude, axis=0)

    x = res_points[:, 0]
    y = res_points[:, 1]

    return x, y
    

test_remove_close_points.py

if __name__ == '__main__':
    import matplotlib.pyplot as plt

    sz = 10000000
    px = np.random.rand(sz) * 10000
    py = np.random.rand(sz) * 10000

    t = time.time()
    fpx, fpy = remove_close_points(px, py, 300)
    print(time.time() - t)

    plt.plot(px, py, '.b', fpx, fpy, 'or')
    plt.axis('equal')
    plt.show()