Vectorizing For Loop With Repeated Indices In Python

I am trying to optimize a snippet that gets called a lot (millions of times) so any type of speed improvement (hopefully removing the for-loop) would be great. I am computing a cor

Solution 1:

Here is the code that use broadcast, hypot, round, bincount to remove all the loops:

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

to compare, I modified your code as follows:

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        

and here is the code to check the results:

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)        

and the %timeit results:

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

outputs:

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

Solution 2:

Your original code on my system runs in about 5.7 seconds. I fully vectorized the inner loop and got it to run in 0.39 seconds. Simply replace your "go through all points" loop with this:

    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 inrange(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                                                                                 

The first observation was that your hypot code was computing 2D distances, so I replaced that with cdist from SciPy to do it all in a single call. The second was that the inner for loop was slow, and thanks to an insightful comment from @hpaulj I vectorized that as well using np.add.at().

---

Since you asked how to vectorize the inner loop as well, I did that later. It now takes 0.25 seconds to run, for a total speedup of over 20x. Here's the final code:

    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

This uses more memory but does produce a substantial speedup vs. the single-loop version above.

Solution 3:

Ok, so as it turns out outer products are incredibly memory expensive, however, using answers from @HYRY and @JohnZwinck i was able to make code that is still roughly linear in n in memory and computes fast (0.5 seconds for the test case)

import numpy as np

deftwopointcorr(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*0for j inrange(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