Как ускорить ближайший поиск в Pandas (возможно, за счет векторизации кода)

У меня есть два фрейма данных. Каждый из них содержит местоположения (X, Y) и значение для этой точки. Для каждой точки в первом кадре данных я хочу найти ближайшую точку во втором кадре данных, а затем найти разницу. У меня есть работающий код, но он использует цикл for, который работает медленно.

Любые предложения о том, как ускорить это? Я знаю, что, как правило, неплохо избавиться от циклов for в пандах для повышения производительности, но я не понимаю, как это сделать в этом случае.

Вот пример кода:

import pandas as pd
import numpy as np

df1=pd.DataFrame(np.random.rand(10,3), columns=['val', 'X', 'Y'])
df2=pd.DataFrame(np.random.rand(10,3), columns=['val', 'X', 'Y'])

nearest=df1.copy()  #CORRECTION.  This had been just =df1 which caused a problem when trying to compare to answers submitted.

for idx,row in nearest.iterrows():
#Find the X,Y points closest to the selected point:
    closest=df2.ix[((df2['X']-row['X'])**2+(df2['Y']-row['Y'])**2).idxmin()]
    #Set the max to the difference between the current row and the nearest one.
    nearest.loc[idx,'val']= df1.loc[idx,'val'] - closest['val'] 

Поскольку я использую это для больших фреймов данных, для вычисления требуется много времени.

Спасибо,


person Adam    schedule 19.02.2015    source источник


Ответы (1)


Одно классное решение вашей проблемы включает использование типа данных complex (встроенного в python и numpy).

import numpy as np
import pandas as pd

df1=pd.DataFrame(np.random.rand(10,3), columns=['val', 'X', 'Y'])
df2=pd.DataFrame(np.random.rand(10,3), columns=['val', 'X', 'Y'])

# dataframes to numpy arrays of complex numbers
p1 = (df1['X'] + 1j * df1['Y']).values
p2 = (df2['X'] + 1j * df2['Y']).values

# calculate all the distances, between each point in
# df1 and each point in df2 (using an array-broadcasting trick)
all_dists = abs(p1[..., np.newaxis] - p2)

# find indices of the minimal distance from df1 to df2,
# and from df2 to df1
nearest_idxs1 = np.argmin(all_dists, axis = 0)
nearest_idxs2 = np.argmin(all_dists, axis = 1)

# extract the rows from the dataframes
nearest_points1 = df1.ix[nearest_idxs1].reset_index()
nearest_points2 = df2.ix[nearest_idxs2].reset_index()

Это, вероятно, намного быстрее, чем использование цикла, но если ваш ряд окажется огромным, он будет потреблять много памяти (квадратичный по количеству точек).

Кроме того, это решение работает, если наборы точек имеют разную длину.


Вот конкретный пример, демонстрирующий, как это работает:

df1 = pd.DataFrame([ [987, 0, 0], [888, 2,2], [2345, 3,3] ], columns=['val', 'X', 'Y'])
df2 = pd.DataFrame([ [ 1000, 1, 1 ], [2000, 9, 9] ] , columns=['val', 'X', 'Y'])

df1
    val  X  Y
0   987  0  0
1   888  2  2
2  2345  3  3

df2
    val  X  Y
0  1000  1  1
1  2000  9  9

Здесь для каждой точки в df1 df2[0]=(1,1) является ближайшей точкой (как показано в nearest_idxs2 ниже). Учитывая обратную задачу, для (1,1) либо (0,0), либо (2,2) являются ближайшими, а для (9,9) df1[1]=(3,3) являются ближайшими ( как показано на nearest_idxs1 ниже).

p1 = (df1['X'] + 1j * df1['Y']).values
p2 = (df2['X'] + 1j * df2['Y']).values
all_dists = abs(p1[..., np.newaxis] - p2)
nearest_idxs1 = np.argmin(all_dists, axis = 0)
nearest_idxs2 = np.argmin(all_dists, axis = 1)

nearest_idxs1
array([0, 2])
nearest_idxs2
array([0, 0, 0])

# It's nearest_points2 you're after:
nearest_points2 = df2.ix[nearest_idxs2].reset_index()

nearest_points2
   index   val  X  Y
0      0  1000  1  1
1      0  1000  1  1
2      0  1000  1  1

df1['val'] - nearest_points2['val']
0     -13
1    -112
2    1345

Чтобы решить обратную задачу (для каждой точки в df2 найти ближайшую в df1), возьмите nearest_points1 и df2['val'] - nearest_points1['val']

person shx2    schedule 19.02.2015
comment
Мне это нравится. Хотя что-то не так получается. Я также попробовал поменять местами оси, как вы предложили. Я считаю, что если я запускаю две версии кода, то ближайшая и ближайшая_точки1 должны быть идентичными. Но это не так. Nearest_points1 (и 2) заканчиваются некоторыми повторяющимися значениями и некоторыми отсутствующими значениями. - person Adam; 19.02.2015
comment
Это связано с тем, что точка в df2 может быть ближайшей к нескольким точкам в df1 или вообще не иметь точек в df1. - person shx2; 19.02.2015
comment
Как может точка в df2 не быть ближайшей точкой в ​​df1? Для любой точки по крайней мере одна точка должна быть ближайшей (их может быть больше, но со случайностями, подобными этому примеру, это маловероятно) - person Adam; 19.02.2015
comment
Да, я согласен. В своем комментарии выше я переключился на роли df1 и df2. Итак: точка в df1 может быть ближайшей к нескольким точкам в df2. - person shx2; 19.02.2015
comment
Я попытался сделать df1 7 строк и df2 5 строк. Поскольку df1 имеет 7 строк, я ожидаю, что для каждой из этих 7 точек он должен найти ближайшую, но вместо этого ближайшая_точка1 имеет только 5 строк. И каждый раз, когда я запускаю его, в ближайших_точках1 появляются повторяющиеся строки (что кажется очень маловероятным для случайных данных). - person Adam; 19.02.2015
comment
Давайте продолжим обсуждение в чате. - person Adam; 19.02.2015
comment
Я не использовал copy(), поэтому мой код изменил df1, и тогда я получил результаты, отличные от вашей версии. Итак, я исправил это, и теперь ваш код дает те же результаты, что и мой. - person Adam; 20.02.2015
comment
Другой интересный результат заключается в том, что для меньших размеров это происходит быстрее. Но для кадров данных большего размера это на самом деле медленнее, чем исходный код. Я думаю, что для хранения больших данных должна использоваться виртуальная память, и это замедляет ее, хотя это работает. - person Adam; 20.02.2015