Векторизация для цикла с повторяющимися индексами в python

Question

Я пытаюсь оптимизировать фрагмент, который называется много (миллионы раз), поэтому любой тип улучшения скорости (надеюсь, удаляющий for-loop) был бы замечательным.

Я вычисления корреляционной функции некоторой частицы -й со всеми остальными

C_j (| г-г '|) = SQRT (E ((s_j (г') - S_K (г))^2)), усредненной по k.

Моя идея состоит в том, чтобы иметь переменную corrfun, которая помещает данные в некоторые ячейки (r, определенные в другом месте). Я нахожу, какой бит из r каждого s_k принадлежит, и он хранится в ind. Таким образом, ind [0] - индекс r (и, следовательно, corrfun), для которого соответствует точка j = 0. Несколько точек могут попадать в один и тот же ящик (на самом деле я хочу, чтобы ячейки были достаточно большими, чтобы содержать несколько точек), поэтому я суммирую все (s_j (r ') - s_k (r))^2, а затем деля по числу точек в этом бункере (хранится в переменной rw). Код, который я в конечном итоге делает для этого является следующее (нп для NumPy):

for k, v in enumerate(ind): 
     if j==k: 
      continue 
     corrfun[v] += (s[k]-s[j])**2 
     rw[v] += 1 
rw2 = rw 
rw2[rw < 1] = 1 
corrfun = np.sqrt(np.divide(corrfun, rw2)) 

записки, RW2 бизнес потому, что я хочу, чтобы избежать деления на 0 проблем, но я вернуть массив Rw и я хочу чтобы иметь возможность различать элементы rw = 0 и rw = 1. Возможно, для этого есть более элегантное решение.

Есть ли способ сделать for-loop быстрее? В то время как я не хотел бы добавлять самостоятельное взаимодействие (j == k), я даже уверен в том, что у меня есть самостоятельное взаимодействие, если это означает, что я могу получить значительно более быстрый расчет (длина ind ~ 1E6, поэтому само взаимодействие, вероятно, будет незначительным).

Спасибо!

Илья

Edit:

Вот полный код. Обратите внимание, что в полном коде я также усредняю ​​по j.

import numpy as np 

def twopointcorr(x,y,s,dr): 

    width = np.max(x)-np.min(x) 
    height = np.max(y)-np.min(y) 

    n = len(x) 

    maxR = np.sqrt((width/2)**2 + (height/2)**2) 

    r = np.arange(0, maxR, dr) 
    print(r) 
    corrfun = r*0 
    rw = r*0 
    print(maxR) 
    ''' go through all points''' 
    for j in range(0, n-1): 
     hypot = np.sqrt((x[j]-x)**2+(y[j]-y)**2) 
     ind = [np.abs(r-h).argmin() for h in hypot] 

     for k, v in enumerate(ind): 
      if j==k: 
       continue 
      corrfun[v] += (s[k]-s[j])**2 
      rw[v] += 1 

    rw2 = rw 
    rw2[rw < 1] = 1 
    corrfun = np.sqrt(np.divide(corrfun, rw2)) 
    return r, corrfun, rw 

отладить проверить это следующим образом

from twopointcorr import twopointcorr 
import numpy as np 
import matplotlib.pyplot as plt 
import time 

n=1000 
x = np.random.rand(n) 
y = np.random.rand(n) 
s = np.random.rand(n) 

print('running two point corr functinon') 

start_time = time.time() 
r,corrfun,rw = twopointcorr(x,y,s,0.1) 
print("--- Execution time is %s seconds ---" % (time.time() - start_time)) 

fig1=plt.figure() 
plt.plot(r, corrfun,'-x') 

fig2=plt.figure() 
plt.plot(r, rw,'-x') 
plt.show() 

Опять же, главный вопрос заключается в том, что в реальных данных п ~ 1E6. Разумеется, я могу переделать его, чтобы уменьшить его, но мне очень хотелось бы на самом деле прокрутить этот набор данных.

Answer 1

Вот код, который используют трансляции, hypot, круглые, bincount удалить все петли:

def twopointcorr2(x, y, s, dr): 
    width = np.max(x)-np.min(x) 
    height = np.max(y)-np.min(y) 
    n = len(x) 
    maxR = np.sqrt((width/2)**2 + (height/2)**2) 
    r = np.arange(0, maxR, dr)  
    osub = lambda x:np.subtract.outer(x, x) 

    ind = np.clip(np.round(np.hypot(osub(x), osub(y))/dr), 0, len(r)-1).astype(int) 
    rw = np.bincount(ind.ravel()) 
    rw[0] -= len(x) 
    corrfun = np.bincount(ind.ravel(), (osub(s)**2).ravel()) 
    return r, corrfun, rw 

для сравнения, я изменил свой код следующим образом:

def twopointcorr(x,y,s,dr): 
    width = np.max(x)-np.min(x) 
    height = np.max(y)-np.min(y) 

    n = len(x) 

    maxR = np.sqrt((width/2)**2 + (height/2)**2) 

    r = np.arange(0, maxR, dr) 
    corrfun = r*0 
    rw = r*0 
    for j in range(0, n): 
     hypot = np.sqrt((x[j]-x)**2+(y[j]-y)**2) 
     ind = [np.abs(r-h).argmin() for h in hypot] 
     for k, v in enumerate(ind): 
      if j==k: 
       continue 
      corrfun[v] += (s[k]-s[j])**2 
      rw[v] += 1 

    return r, corrfun, rw   

и здесь код для проверки результатов:

import numpy as np 

n=1000 
x = np.random.rand(n) 
y = np.random.rand(n) 
s = np.random.rand(n) 

r1, corrfun1, rw1 = twopointcorr(x,y,s,0.1) 
r2, corrfun2, rw2 = twopointcorr2(x,y,s,0.1) 

assert np.allclose(r1, r2) 
assert np.allclose(corrfun1, corrfun2) 
assert np.allclose(rw1, rw2)   

и результаты% timeit:

%timeit twopointcorr(x,y,s,0.1) 
%timeit twopointcorr2(x,y,s,0.1)  

выходов:

1 loop, best of 3: 5.16 s per loop 
10 loops, best of 3: 134 ms per loop 

Answer 2

Ваш оригинальный код в моей системе длится около 5,7 секунд. Я полностью векторизовал внутренний цикл и запустил его за 0,39 секунды. Просто замените ваш «пройти через все точки» цикл с этим:

points = np.column_stack((x,y)) 
    hypots = scipy.spatial.distance.cdist(points, points) 
    inds = np.rint(hypots.clip(max=maxR)/dr).astype(np.int) 

    # go through all points    
    for j in range(n): # n.b. previously n-1, not sure why 
     ind = inds[j] 

     np.add.at(corrfun, ind, (s - s[j])**2) 

     np.add.at(rw, ind, 1) 
     rw[ind[j]] -= 1 # subtract self                     

Первое наблюдение было то, что ваш hypot код был вычисления 2D расстояния, поэтому я заменил, что с cdist от SciPy, чтобы сделать все это в одном вызове. Во-вторых, внутренняя петля for была медленной, и благодаря проницательному комментарию от @hpaulj I, векторизованного так же, как и с использованием np.add.at().

Поскольку вы спросили, как векторизовать внутренний цикл, я сделал это позже. Теперь требуется 0.25 секунд для полного ускорения более 20 раз. Вот окончательный код:

points = np.column_stack((x,y)) 
    hypots = scipy.spatial.distance.cdist(points, points) 
    inds = np.rint(hypots.clip(max=maxR)/dr).astype(np.int) 

    sn = np.tile(s, (n,1)) # n copies of s                    
    diffs = (sn - sn.T)**2 # squares of pairwise differences 
    np.add.at(corrfun, inds, diffs) 

    rw = np.bincount(inds.flatten(), minlength=len(r)) 
    np.subtract.at(rw, inds.diagonal(), 1) # subtract self 

Это использует больше памяти, но производит существенное ускорение против.одноконтурная версия выше.

Answer 3

Ok, так как получается внешние продукты невероятно память дорого, однако, используя ответы из @HYRY и @JohnZwinck я был в состоянии сделать код, который до сих пор примерно линейным in n в памяти и быстро вычисляется (0,5 с для тестового случая)

import numpy as np 

def twopointcorr(x,y,s,dr,maxR=-1): 

    width = np.max(x)-np.min(x) 
    height = np.max(y)-np.min(y) 

    n = len(x) 

    if maxR < dr: 
     maxR = np.sqrt((width/2)**2 + (height/2)**2) 

    r = np.arange(0, maxR+dr, dr) 

    corrfun = r*0 
    rw = r*0 


    for j in range(0, n): 

     ind = np.clip(np.round(np.hypot(x[j]-x,y[j]-y)/dr), 0, len(r)-1).astype(int) 
     np.add.at(corrfun, ind, (s - s[j])**2) 
     np.add.at(rw, ind, 1) 

    rw[0] -= n 

    corrfun = np.sqrt(np.divide(corrfun, np.maximum(rw,1))) 
    r=np.delete(r,-1) 
    rw=np.delete(rw,-1) 
    corrfun=np.delete(corrfun,-1) 
    return r, corrfun, rw