高效地查找多个维度上的邻居并根据邻近度计算值的总和

Question

我的任务是找出中心元素可变距离内所有元素的总值。元素使用 3 个维度（我的数据中的列）排列。每个元素都有一个给定 3 个维度的唯一位置（并且有一个唯一 ID）。

我有一个可以满足我要求的工作版本，但是速度非常慢。我正在使用 itertuples，使用子集数据框 apply(np.isclose) 查找每个元组的值，然后使用 .at 设置值（请参见下面的代码）。

问题不在于我的代码的功能，而在于可扩展性。由于我想设置一个可变距离来测量，并且我想为每一行计算这个值，它最终迭代 nrows x ndistances，目前每次迭代需要 1.7 秒（我的数据有 >25,000 行，我估计 ~12 小时我尝试的每一个距离）。

import pandas as pd
import numpy as np

数据结构示例：

df = pd.DataFrame({'id':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19], 
                          'x':[-2,-2,-2,-1,-1,-1,-1,0,0,0,0,0,1,1,1,1,2,2,2], 
                          'y':[2,1,0,2,1,0,-1,2,1,0,-1,-2,1,0,-1,-2,0,-1,-2], 
                          'z':[0,1,2,-1,0,1,2,-2,-1,0,1,2,-2,-1,0,1,-2,-1,0], 
                          'val':[0,0,0,1,0,0,6,3,7,11,0,0,14,18,10,4,20,15,2]})
df.set_index('id', inplace=True)
# The 'val' column can have any non-negative whole number, I've just picked some randomly.

到目前为止的“工作”代码：

n = 0  #Initial distance
while n < 3:  #This part allows me to set my distance range
    df['n{0}'.format(n)] = np.nan  #create a column for the new values
    for row in df.itertuples():
        valsum = df[(df['x'].apply(np.isclose, b=row.x, atol=n)) & 
                    (df['y'].apply(np.isclose, b=row.y, atol=n)) & 
                    (df['z'].apply(np.isclose, b=row.z, atol=n))].val.sum()
        df.at[row.Index, 'n{0}'.format(n)] = valsum
    n += 1

当前/期望的输出：

    x   y   z   val n0  n1  n2
id                          
1   -2  2   0   0   0   1   22
2   -2  1   1   0   0   0   25
3   -2  0   2   0   0   6   17
4   -1  2   -1  1   1   11  54
5   -1  1   0   0   0   19  70
6   -1  0   1   0   0   17  57
7   -1  -1  2   6   6   6   31
8   0   2   -2  3   3   25  74
9   0   1   -1  7   7   54  99
10  0   0   0   11  11  46  111
11  0   -1  1   0   0   31  73
12  0   -2  2   0   0   10  33
13  1   1   -2  14  14  62  99
14  1   0   -1  18  18  95  105
15  1   -1  0   10  10  60  107
16  1   -2  1   4   4   16  66
17  2   0   -2  20  20  67  100
18  2   -1  -1  15  15  65  101
19  2   -2  0   2   2   31  80

我知道“n0”列等于“val”列，因为搜索距离为 0，但我希望显示我正在寻找的内容。 val列中所有项的总和为111，当(x,y,z) = (0,0,0)时相同。这是因为 (0,0,0) 是本示例中我的数据的中心，因此距离为 2 会捕获所有元素。我想在一定距离的带宽内执行此操作，例如 5-10。

我的终极问题是：我怎样才能更快/更有效地做到这一点？

Answer 1

此解决方案还使用 KDTrees（来自 scipy 库）。

在您的代码和前面的答案中，当循环计算半径 = 3 的结果时，它正在重复半径 = 0、1 已经完成的工作， 和 2.

下面的代码通过节点一次完成所有计算。定义最大距离和范围箱数。查找具有最大距离的所有节点对并使用np.digitize() 将实际距离映射到范围箱。将 'val' 添加到映射范围 bin。

import pandas as pd
import numpy as np

from scipy.spatial import cKDTree as KDTree

# define the range and number of range bins 
# this example defines 3 bins: 0.0 - 1.0; 1.0 - 2.0; 2.0 - 3.0
max_distance = 3.0
nbins = 3
bin_range = 0.0, max_distance
bins = np.linspace(*bin_range, nbins+1)[1:]

# build a KDTree and generate a sparse matrix of node pairs
# that have a max distance of bin_range[-1]
tree = KDTree(df[['x','y','z']])
dist = tree.sparse_distance_matrix(tree, bin_range[-1])

# one row per node, one column per range bin
sums = np.zeros((len(df), nbins))

# for each pair of nodes, map the range to the bin index and add
# the value of the second node to mapped bin for the 1st node 
for (j,k),d in dist.items():
    sums[j][np.digitize(d, bins)] += df['val'][k+1]

对于每个节点，数组sums 包含一行，其中包含分箱范围的总和。例如，第一列包含距离

sums

array([[ 0.,  1., 21.],
       [ 0.,  0., 25.],
       [ 0.,  6., 11.],
       [ 1., 10., 43.],
       [ 0., 19., 51.],
       [ 0., 17., 40.],
       [ 6.,  0., 25.],
       [ 3., 22., 49.],
       [ 7., 47., 45.],
       [11., 35., 65.],
       [ 0., 31., 42.],
       [ 0., 10., 23.],
       [14., 48., 37.],
       [18., 77., 10.],
       [10., 50., 47.],
       [ 4., 12., 50.],
       [20., 47., 33.],
       [15., 50., 36.],
       [ 2., 29., 49.]])

Answer 2

这是一个不需要额外软件包的解决方案。

这些函数定义了两点a 和b 之间的距离。这里显示了欧几里得距离、曼哈顿距离和切比雪夫距离（感谢@Peter Leimbigler answer，他认识到最后一个是 OP 使用的距离）。 a 和 b 被假定为 3 长度列表。您可以使用其中之一（甚至可以定义其他自定义距离函数）。

def euclidean(a, b):
    """euclidean distance"""
    return np.sqrt((a[0] - b[0])**2 + (a[1] - b[1])**2 + (a[2] - b[2])**2) 

def manhattan(a, b):
    """manhattan distance"""
    return abs(a[0] - b[0]) + abs(a[1] - b[1]) + abs(a[2] - b[2])

def cebyshev(a, b):
    """cebyshev distance"""
    return max(abs(a[0] - b[0]), abs(a[1] - b[1]), abs(a[2] - b[2]))

以下函数为点point 返回数据框data（这是您的数据框）中val 列的值的总和，其坐标小于距离d。 func 是用于计算距离的函数（之前的函数之一）。

def getclosesum(data, point, d, func):
    dists = data.apply(lambda x : func(x, point), axis=1)
    return data['val'].loc[dists <= d].sum()

最后你可以使用df.apply来计算你的列：

for n in range(3):
    df['n{0}'.format(n)] = df.apply(lambda x : getclosesum(df, x, n, cebyshev), axis=1)

使用您的示例数据框，在我的机器上，此代码需要 0.155 秒才能完成这项工作，而您的原始代码需要 0.233 秒。
所以这比您的解决方案要快，但没有@Peter Leimbigler 提供的代码那么快（我敢打赌scikit 更加优化）。

Answer 3

在 k 维空间内寻找最近邻是 k-d 树数据结构 (Wikipedia) 的经典案例。 Scikit-learn 有一个灵活的实现 (docs)，我在下面使用它，因为您问题中使用的条件逻辑似乎定义了 scikit-learn 本身支持的 Chebyshev 距离度量 (Wikipedia)。 SciPy 的 cKDTree (docs, C++ source code) 仅支持欧几里得 (L2) 距离度量，但针对它进行了优化，因此可能更快。

# Setup
df = pd.DataFrame({'id':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19], 
                   'x':[-2,-2,-2,-1,-1,-1,-1,0,0,0,0,0,1,1,1,1,2,2,2], 
                   'y':[2,1,0,2,1,0,-1,2,1,0,-1,-2,1,0,-1,-2,0,-1,-2], 
                   'z':[0,1,2,-1,0,1,2,-2,-1,0,1,2,-2,-1,0,1,-2,-1,0], 
                   'val':[0,0,0,1,0,0,6,3,7,11,0,0,14,18,10,4,20,15,2]})
df.set_index('id', inplace=True)


from sklearn.neighbors import KDTree

# Build k-d tree with the Chebyshev metric, AKA L-infinity
tree = KDTree(df[['x', 'y', 'z']].values, metric='chebyshev')

for radius in [0, 1, 2]:
    # Populate new column with placeholder integer
    df[f'n{radius}'] = -1
    for i, row in df.iterrows():
        coords = row[['x', 'y', 'z']].values.reshape(1, -1)
        idx = tree.query_radius(coords, r=radius)[0]
        df.loc[i, f'n{radius}'] = df.iloc[idx]['val'].sum()

df
    x  y  z  val  n0  n1   n2
id                           
1  -2  2  0    0   0   1   22
2  -2  1  1    0   0   0   25
3  -2  0  2    0   0   6   17
4  -1  2 -1    1   1  11   54
5  -1  1  0    0   0  19   70
6  -1  0  1    0   0  17   57
7  -1 -1  2    6   6   6   31
8   0  2 -2    3   3  25   74
9   0  1 -1    7   7  54   99
10  0  0  0   11  11  46  111
11  0 -1  1    0   0  31   73
12  0 -2  2    0   0  10   33
13  1  1 -2   14  14  62   99
14  1  0 -1   18  18  95  105
15  1 -1  0   10  10  60  107
16  1 -2  1    4   4  16   66
17  2  0 -2   20  20  67  100
18  2 -1 -1   15  15  65  101
19  2 -2  0    2   2  31   80