Блог

Я ищу ближайшую ненулевую ячейку в трехмерном массиве numpy на основе координат i,j,k, хранящихся во фрейме данных pandas. Мое решение ниже работает, но оно медленнее, чем мне бы хотелось. Я знаю, что мне не хватает навыков оптимизации, поэтому я надеюсь, что кто-нибудь сможет помочь мне найти более быстрый вариант.

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

 a=np.random.randint(0,2,size=(100,100,100))
# df with i,j,k of interest
df=pd.DataFrame(np.random.randint(100,size=(100,3)).tolist(),
                columns=['i','j','k'])

def find_nearest(a,df):
    import numpy as np
    import pandas as pd
    import time
    t0=time.time()
    nzi = np.nonzero(a)
    for i,r in df.iterrows():
        dist = ((r['k'] - nzi[0])**2   
              (r['i'] - nzi[1])**2   
              (r['j'] - nzi[2])**2)
        nidx = dist.argmin()
        df.loc[i,['nk','ni','nj']]=(nzi[0][nidx],
                                    nzi[1][nidx],
                                    nzi[2][nidx])
    print(time.time()-t0)
    return(df)
 

Проблема, которую вы пытаетесь решить, выглядит как поиск ближайшего соседа.

Наихудшая сложность текущего кода O(n m) связана с n количеством точек для поиска и m количеством соседних кандидатов. С n = 100 и m = 100**3 = 1,000,000 это означает около сотен миллионов итераций. Чтобы эффективно решить эту проблему, можно использовать более совершенный алгоритм.

Распространенный способ решения такого рода проблем состоит в размещении всех элементов в структуре данных BSP-дерева (например, Квадродерево или Октрево. Такая структура данных помогает вам находить ближайшие элементы в определенном месте за определенный O(log(m)) промежуток времени. В результате общая сложность этого метода составляет O(n log(m)) ! SciPy уже реализует KD-деревья.

Векторизация, как правило, также помогает ускорить вычисления.

 def find_nearest_fast(a,df):
    from scipy.spatial import KDTree
    import numpy as np
    import pandas as pd
    import time
    t0=time.time()
    candidates = np.array(np.nonzero(a)).transpose().copy()
    tree = KDTree(candidates, leafsize=1024, compact_nodes=False)
    searched = np.array([df['k'], df['i'], df['j']]).transpose()
    distances, indices = tree.query(searched)
    nearestPoints = candidates[indices,:]
    df[['nk', 'ni', 'nj']] = nearestPoints
    print(time.time()-t0)
    return df
 

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