Vectorizing for loop with repeated indices in python
Question: 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 correlation function of some j'th particle with all others

C_j(|r-r'|) = sqrt(E((s_j(r')-s_k(r))^2)) averaged over k.

My idea is to have a variable corrfun which bins data into some bins (the r,
defined elsewhere). I find what bin of r each s_k belongs to and this is
stored in ind. So ind[0] is the index of r (and thus the corrfun) for which
the j=0 point corresponds to. Multiple points can fall into the same bin (in
fact I want bins to be big enough to contain multiple points) so I sum
together all of the (s_j(r')-s_k(r))^2 and then divide by number of points in
that bin (stored in variable rw). The code I ended up making for this is the
following (np is for 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))
    

Note, the rw2 business was because I want to avoid divide by 0 problems but I
do return the rw array and I want to be able to differentiate between the rw=0
and rw=1 elements. Perhaps there is a more elegant solution for this as well.

Is there a way to make the for-loop faster? While I would like to not add the
self interaction (j==k) I am even ok with having self interaction if it means
I can get significantly faster calculation (length of ind ~ 1E6 so self
interaction is probably insignificant anyways).

Thank you!

Ilya

Edit:

Here is the full code. Note, in the full code I am averaging over j as well.

    
    
    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
    

I debug test it the following way

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

Again, the main issue is that in the real dataset n~1E6. I can resample to
make it smaller, of course, but I would love to actually crank through the
dataset.


Answer: 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 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                                                                                 
    

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.

