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

我的任务是找到中心元素可变距离内所有元素的总价值。这些元素使用 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）内执行此操作。

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

在 k 维空间内查找最近邻居是 k-d 树数据结构的经典案例（维基百科 https://en.wikipedia.org/wiki/K-d_tree）。 Scikit-learn 有一个灵活的实现（docs https://scikit-learn.org/stable/modules/generated/sklearn.neighbors.KDTree.html）我在下面使用它，因为你的问题中使用的条件逻辑似乎定义了切比雪夫距离度量（维基百科 https://en.wikipedia.org/wiki/Chebyshev_distance），scikit-learn 本身支持。 SciPy 的cKDTree (docs https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.cKDTree.html, C++源代码 https://github.com/scipy/scipy/tree/master/scipy/spatial/ckdtree/src) 仅支持欧几里德 (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