加快熊猫DataFrame的嵌套循环

时间:2019-12-22 19:52:46

标签: python pandas dataframe

我有一个pandas.DataFrame,其中包含许多按identityxy排列的对象的坐标。

我正在尝试找到两个身份最接近的对象。要弄清楚我的意思,请使用以下代码:

import numpy as np
import pandas as pd

# Generate random data
df_identity_1 = pd.DataFrame({'identity':1, 'x':np.random.randn(10000), 'y':np.random.randn(10000)})
df_identity_2 = pd.DataFrame({'identity':2, 'x':np.random.randn(10000), 'y':np.random.randn(10000)})
df = pd.concat([df_identity_1, df_identity_2])

>>> df
      identity         x         y
0            1 -1.784748  2.085517
1            1  0.324645 -1.584790
2            1 -0.044623 -0.348576
3            1  0.802035  1.362336
4            1 -0.091508 -0.655114
...        ...       ...       ...
9995         2  0.939491  0.304964
9996         2 -0.233707 -0.135265
9997         2  0.792494  1.157236
9998         2 -0.385080 -0.021226
9999         2  0.105970 -0.042135

当前,我必须遍历每一行并再次遍历整个DataFrame以找到最接近的坐标。

# Function to find the absolute / Euclidean distance between two coordinates
def euclidean(x1, y1, x2, y2):
    a = np.array((int(x1), int(y1)))
    b = np.array((int(x2), int(y2)))
    return np.linalg.norm(a-b)

# Function to find the closest coordinate with a different index
def find_closest_coord(row, df):
    d = df[(df['identity'] != int(row['identity']))]
    if d.empty:
        return None
    return min(euclidean(row.x, row.y, r.x, r.y) for r in df.itertuples(index=False))

df['closest_coord'] = df.apply(lambda row: find_closest_coord(row, df), axis=1)

此代码具有完整的功能-但是当我有一个大型数据集(+ 100k坐标)时,此“嵌套” for循环非常耗时。

是否存在某些功能可以加快这一概念或更快地实现这一目标?

1 个答案:

答案 0 :(得分:1)

解决此问题的最佳方法是使用空间数据结构。这些数据结构使您可以在需要进行此类查询时大大减少搜索空间的大小。 SciPy为最近的邻居查询提供了KD树,但是将其扩展到多台机器(如果您的数据大小需要这样做)会很麻烦。

如果您需要扩展规模,则可能需要使用专用的地理空间分析工具。

通常,如果要加快这样的速度,则需要在迭代方法和内存强度之间进行权衡。

但是,在这种情况下,您的核心瓶颈是:

  • 逐行迭代
  • 每两行调用一次euclidean ,而不是每个数据集调用一次

诸如norm之类的NumPy函数本质上是列式的,您应该通过在整个数据数组上调用它来利用它。如果每个数据框都是10,000行,则您要调用norm 1亿次。稍微调整一下以进行更改将对您有很大帮助。

如果要在Python中大规模执行此操作,并且不能有效地使用空间数据结构(并且不想使用启发式方法来减少搜索空间),则可能会执行以下操作:跨产品合并两个表,一次进行一次列运算即可计算出欧氏距离一次,然后使用groupby-aggregation(min)来获取最接近的点。

与您逐行进行迭代相比,这将更快并且要占用更多的内存,但是可以使用Dask(或Spark)之类的方法轻松扩展。

我将仅用几行来说明逻辑。

import numpy as np
import pandas as pd

# Generate random data
nrows = 3
df_identity_1 = pd.DataFrame({'identity':1, 'x':np.random.randn(nrows), 'y':np.random.randn(nrows)})
df_identity_2 = pd.DataFrame({'identity':2, 'x':np.random.randn(nrows), 'y':np.random.randn(nrows)})
df_identity_1.reset_index(drop=False, inplace=True)
df_identity_2.reset_index(drop=False, inplace=True)

除了每个数据帧的identity标志之外,请注意如何创建唯一索引。稍后将对groupby派上用场。接下来,我可以进行跨产品联接。如果我们使用不同的列名,这会更干净,但我将使其与您的示例保持一致。随着数据集的增长,这种联接将很快在纯熊猫中耗尽内存,但是Dask(https://dask.org/)将能够很好地处理它。

def cross_product(left, right):
    return left.assign(key=1).merge(right.assign(key=1), on='key').drop('key', 1)

crossprod = cross_product(df_identity_1, df_identity_2)
crossprod
index_x identity_x  x_x y_x index_y identity_y  x_y y_y
0   0   1   1.660468    -1.954339   0   2   -0.431543   0.500864
1   0   1   1.660468    -1.954339   1   2   -0.607647   -0.436480
2   0   1   1.660468    -1.954339   2   2   1.613126    -0.696860
3   1   1   0.153419    0.619493    0   2   -0.431543   0.500864
4   1   1   0.153419    0.619493    1   2   -0.607647   -0.436480
5   1   1   0.153419    0.619493    2   2   1.613126    -0.696860
6   2   1   -0.592440   -0.299046   0   2   -0.431543   0.500864
7   2   1   -0.592440   -0.299046   1   2   -0.607647   -0.436480
8   2   1   -0.592440   -0.299046   2   2   1.613126    -0.696860

接下来,我们只需要计算每行的最小距离,然后分别按index_xindex_y分组,即可获得最小距离值。请注意,我们如何通过一次调用norm来实现此目的,而不是每行一次。

crossprod['dist'] = np.linalg.norm(crossprod[['x_x', 'y_x']].values - crossprod[['x_y', 'y_y']].values, axis=1)
closest_per_identity1 = crossprod.groupby(['index_x']).agg({'dist':'min'})
closest_per_identity2 = crossprod.groupby(['index_y']).agg({'dist':'min'})
closest_per_identity1
dist
index_x 
0   1.258370
1   0.596869
2   0.138273
closest_per_identity2
dist
index_y 
0   0.596869
1   0.138273
2   1.258370

与您在相同数据上的原始示例相比。请注意,我将您的int调用更改为floats,并且您的迭代遍历了d,而不是df(否则,您是将每个点与其自身进行比较)。 / p>

df = pd.concat([df_identity_1, df_identity_2])
​
def euclidean(x1, y1, x2, y2):
    a = np.array((float(x1), float(y1)))
    b = np.array((float(x2), float(y2)))
    return np.linalg.norm(a-b)
​
# Function to find the closest coordinate with a different index
def find_closest_coord(row, df):
    d = df[(df['identity'] != int(row['identity']))]
    if d.empty:
        return None
    r = min(euclidean(row.x, row.y, r.x, r.y) for r in d.itertuples(index=False))
    return r
​
df['closest_coord'] = df.apply(lambda row: find_closest_coord(row, df), axis=1)
df
index   identity    x   y   closest_coord
0   0   1   1.660468    -1.954339   1.258370
1   1   1   0.153419    0.619493    0.596869
2   2   1   -0.592440   -0.299046   0.138273
0   0   2   -0.431543   0.500864    0.596869
1   1   2   -0.607647   -0.436480   0.138273
2   2   2   1.613126    -0.696860   1.258370