查找所有比给定最大距离近的点对

Question

我想（有效地）找到所有比某个距离更近的点对max_d。我当前使用cdist的方法是：

import numpy as np
from scipy.spatial.distance import cdist

def close_pairs(X,max_d):
    d = cdist(X,X)

    I,J = (d<max_d).nonzero()
    IJ  = np.sort(np.vstack((I,J)), axis=0)

    # remove diagonal element
    IJ  = IJ[:,np.diff(IJ,axis=0).ravel()<>0]

    # remove duplicate
    dt = np.dtype([('i',int),('j',int)])
    pairs = np.unique(IJ.T.view(dtype=dt)).view(int).reshape(-1,2)

    return pairs

def test():
    X = np.random.rand(100,2)*20
    p = close_pairs(X,2)

    from matplotlib import pyplot as plt
    plt.clf()
    plt.plot(X[:,0],X[:,1],'.r')
    plt.plot(X[p,0].T,X[p,1].T,'-b')

但我认为这是一个过度的杀戮（而且不太可读），因为大部分的工作只是为了消除与自我和重复的距离。在

我的主要问题是：有更好的方法吗？在

（注意：输出类型（array，set，…）在这一点上并不重要）

{My-thinking数组{My-thinking-on-current}只返回当前的距离。然而，一旦我从压缩距离数组中找到合适的坐标k，我如何计算它与哪个i,j对等价？在

所以另一个问题是：有没有一种简单的方法来获得pdist输出项的坐标对列表：

• 函数f(k)->i,j
• 使cdist(X,X)[i,j] = pdist(X)[k]

Answer 1

<>根据我的经验，在3D中找到邻居列表有两种最快的方法：一种是使用最简单的双环来编写C++或Cython的循环代码（在我的例子中，两者都是）。它在N^2下运行，但对于小型系统来说非常快。另一种方法是使用线性时间算法。Scipy ckdtree是一个不错的选择，但也有局限性。分子动力学软件中的邻居列表查找器功能最强大，但很难包装，而且初始化时间可能很慢。在

下面我比较四种方法：

• 天真的cython代码
• 围绕OpenMM的包装器（很难安装，请参见下文）
• Scipy.spatial.ckdtree
• scipy.spatial.distance.pdist

测试设置：n点分散在体积密度为0.2的矩形框中。系统大小从10到1000000（一百万）个粒子。接触半径取自0.5, 1, 2, 4, 7, 10。注意，因为密度是0.2，在接触半径为0.5时，我们平均每个粒子有0.1个接触点，1=0.8，2=6.4，10-大约800！对小系统重复多次接触发现，对大于30k粒子的系统重复一次。如果每次呼叫的时间超过5秒，则中止运行。在

设置：dual xeon 2687Wv3，128GB内存，Ubuntu 14.04，python 2.7.11，scipy 0.16.0，numpy 1.10.1。所有的代码都没有使用并行优化（除了OpenMM，尽管并行部分运行得非常快，以至于在CPU图上甚至都不明显，但大部分时间都是通过管道将数据从OpenMM传递出去）。在

结果：请注意，下面的曲线图是对数标度，分布在6个数量级上。即使是很小的视觉差异，实际上也可能是10倍。 对于少于1000个粒子的系统，Cython代码总是更快。但是，1000个粒子之后，结果取决于接触半径。pdist实现总是比cython慢，并且占用更多的内存，因为它显式地创建了一个距离矩阵，这是由于sqrt而导致的。在

• 在较小的接触半径下（每个粒子1个接触），ckdtree是所有系统尺寸的好选择。在
• 在中等接触半径（每个粒子5-50个接触）下，朴素的cython实现是最好的，最多10000个粒子，然后OpenMM开始以几个数量级的优势获胜，但是ckdtree的性能仅差3-10倍
• 在高接触半径（每个粒子接触200个以上）时，简单的方法可以处理多达100k或1M的粒子，然后OpenMM可能会获胜。在

安装OpenMM非常棘手；您可以在http://bitbucket.org/mirnylab/openmm-polymer文件中阅读更多内容”联系人地图.py“或者在自述中。然而，下面的结果表明，对于N>；100k粒子，每个粒子只有5-50个接触点才是有利的。在

Cython代码如下：

import numpy as np
cimport numpy as np
cimport cython

cdef extern from "<vector>" namespace "std":
    cdef cppclass vector[T]:
        cppclass iterator:
            T operator*()
            iterator operator++()
            bint operator==(iterator)
            bint operator!=(iterator)
        vector()
        void push_back(T&)
        T& operator[](int)
        T& at(int)
        iterator begin()
        iterator end()

np.import_array() # initialize C API to call PyArray_SimpleNewFromData
cdef public api tonumpyarray(int* data, long long size) with gil:
    if not (data and size >= 0): raise ValueError
    cdef np.npy_intp dims = size
    #NOTE: it doesn't take ownership of `data`. You must free `data` yourself
    return np.PyArray_SimpleNewFromData(1, &dims, np.NPY_INT, <void*>data)

@cython.boundscheck(False)
@cython.wraparound(False)
def contactsCython(inArray, cutoff):
    inArray = np.asarray(inArray, dtype = np.float64, order = "C")
    cdef int N = len(inArray)
    cdef np.ndarray[np.double_t, ndim = 2] data = inArray
    cdef int j,i
    cdef double curdist
    cdef double cutoff2 = cutoff * cutoff  # IMPORTANT to avoid slow sqrt calculation
    cdef vector[int] contacts1
    cdef vector[int] contacts2
    for i in range(N):
        for j in range(i+1, N):
            curdist = (data[i,0] - data[j,0]) **2 +(data[i,1] - data[j,1]) **2 + (data[i,2] - data[j,2]) **2
            if curdist < cutoff2:
                contacts1.push_back(i)
                contacts2.push_back(j)
    cdef int M = len(contacts1)

    cdef np.ndarray[np.int32_t, ndim = 2] contacts = np.zeros((M,2), dtype = np.int32)
    for i in range(M):
        contacts[i,0] = contacts1[i]
        contacts[i,1] = contacts2[i]
    return contacts

Cython代码的编译（或生成文件）：

^{pr2}$

测试代码：

from __future__ import print_function, division

import signal
import time
from contextlib import contextmanager

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial import ckdtree
from scipy.spatial.distance import pdist

from contactmaps import giveContactsOpenMM  # remove this unless you have OpenMM and openmm-polymer libraries installed
from fastContacts import contactsCython


class TimeoutException(Exception): pass


@contextmanager
def time_limit(seconds):
    def signal_handler(signum, frame):
        raise TimeoutException("Timed out!")

    signal.signal(signal.SIGALRM, signal_handler)
    signal.alarm(seconds)
    try:
        yield
    finally:
        signal.alarm(0)


matplotlib.rcParams.update({'font.size': 8})


def close_pairs_ckdtree(X, max_d):
    tree = ckdtree.cKDTree(X)
    pairs = tree.query_pairs(max_d)
    return np.array(list(pairs))


def condensed_to_pair_indices(n, k):
    x = n - (4. * n ** 2 - 4 * n - 8 * k + 1) ** .5 / 2 - .5
    i = x.astype(int)
    j = k + i * (i + 3 - 2 * n) / 2 + 1
    return np.array([i, j]).T


def close_pairs_pdist(X, max_d):
    d = pdist(X)
    k = (d < max_d).nonzero()[0]
    return condensed_to_pair_indices(X.shape[0], k)


a = np.random.random((100, 3)) * 3  # test set
methods = {"cython": contactsCython, "ckdtree": close_pairs_ckdtree, "OpenMM": giveContactsOpenMM,
           "pdist": close_pairs_pdist}

# checking that each method gives the same value
allUniqueInds = []
for ind, method in methods.items():
    contacts = method(a, 1)
    uniqueInds = contacts[:, 0] + 100 * contacts[:, 1]  # unique index of each contacts
    allUniqueInds.append(np.sort(uniqueInds))  # adding sorted unique conatcts
for j in allUniqueInds:
    assert np.allclose(j, allUniqueInds[0])

# now actually doing testing
repeats = [30,30,30, 30, 30, 20,  20,   10,   5,   3,     2 ,       1,     1,      1]
sizes =    [10,30,100, 200, 300,  500, 1000, 2000, 3000, 10000, 30000, 100000, 300000, 1000000]
systems = [[np.random.random((n, 3)) * ((n / 0.2) ** 0.333333) for k in range(repeat)] for n, repeat in
           zip(sizes, repeats)]

for j, radius in enumerate([0.5, 1, 2, 4, 7, 10]):
    plt.subplot(2, 3, j + 1)
    plt.title("Radius = {0}; {1:.2f} cont per particle".format(radius, 0.2 * (4 / 3 * np.pi * radius ** 3)))

    times = {i: [] for i in methods}

    for name, method in methods.items():
        for n, system, repeat in zip(sizes, systems, repeats):
            if name == "pdist" and n > 30000:
                break  # memory issues
            st = time.time()
            try:
                with time_limit(5 * repeat):
                    for ind in range(repeat):
                        k = len(method(system[ind], radius))
            except:
                print("Run aborted")
                break
            end = time.time()
            mytime = (end - st) / repeat
            times[name].append((n, mytime))
            print("{0} radius={1} n={2} time={3} repeat={4} contPerParticle={5}".format(name, radius, n, mytime,repeat, 2 * k / n))

    for name in sorted(times.keys()):
        plt.plot(*zip(*times[name]), label=name)
    plt.xscale("log")
    plt.yscale("log")
    plt.xlabel("System size")
    plt.ylabel("Time (seconds)")
    plt.legend(loc=0)

plt.show()

Answer 2

我终于自己找到了。将压缩距离数组中的索引k转换为平方距离数组中的等价i,j的函数为：

def condensed_to_pair_indices(n,k):
    x = n-(4.*n**2-4*n-8*k+1)**.5/2-.5
    i = x.astype(int)
    j = k+i*(i+3-2*n)/2+1
    return i,j

我不得不和sympy玩了一会儿才能找到它。现在，要计算相距小于给定距离的所有点对：

^{pr2}$

正如预期的那样，它比其他方法更有效（但是我没有测试ckdtree）。我将更新timeit答案。在

Answer 3

下面是如何使用cKDTree模块来实现这一点。见query_pairs

import numpy as np
from scipy.spatial.distance import cdist
from scipy.spatial import ckdtree


def close_pairs(X,max_d):
    d = cdist(X,X)

    I,J = (d<max_d).nonzero()
    IJ  = np.sort(np.vstack((I,J)), axis=0)

    # remove diagonal element
    IJ  = IJ[:,np.diff(IJ,axis=0).ravel()<>0]

    # remove duplicate
    dt = np.dtype([('i',int),('j',int)])
    pairs = np.unique(IJ.T.view(dtype=dt)).view(int).reshape(-1,2)

    return pairs

def close_pairs_ckdtree(X, max_d):
    tree = ckdtree.cKDTree(X)
    pairs = tree.query_pairs(max_d)
    return np.array(list(pairs))

def test():
    np.random.seed(0)
    X = np.random.rand(100,2)*20
    p = close_pairs(X,2)
    q = close_pairs_ckdtree(X, 2)

    from matplotlib import pyplot as plt
    plt.plot(X[:,0],X[:,1],'.r')
    plt.plot(X[p,0].T,X[p,1].T,'-b')
    plt.figure()
    plt.plot(X[:,0],X[:,1],'.r')
    plt.plot(X[q,0].T,X[q,1].T,'-b')

    plt.show()

t