Numpy vectorization for nearest neighbor calculation in semi-sparse arrays
<p>I have a question for you guys concerning vectorization with numpy.</p>

<p>I have a semi-sparse matrix where the non-zero elements are either +1 or –1 and are ordered randomly in a numpy array (let's call lattice) such that a 'site' at lattice[i,j] has nearest neighbors at lattice[i+1,j+1], lattice[i+1,j-1], lattice[i-1,j+1], and lattice[i-1,j-1]:</p>

<pre><code>lattice = np.array([[1,  0, -1,  0],
                    [0,  1,  0, -1],
                    [1,  0,  1,  0],
                    [0,  1,  0, -1]]) 
</code></pre>

<p>There's a function F that maps the configuration stored in lattice to a scalar value. Not sure how to display mathjax / latex here but I'll do my best with typing it out.</p>

<p>F(lattice) =  a * sum( lattice[i,j] ) + b * sum( lattice[i,j] * lattice[i±1, j±1] )</p>

<p>The first expression is a sum over all of the lattice sites which are scaled by a, and the second sum is over all of the lattice sites and their product with each site's set of nearest neighbors, scaled by b.</p>

<p>In an effort to simplify all of this, when I generate my lattice, I create a dictionary for which each key corresponds to a tuple of indices for each site in lattice, and the value stored there is a list of tuples pointing to the nearest neighbors, i.e. </p>

<pre><code>sites = { (i,j) : [ (i+1, j+1), (i+1, j-1), (i-1,j+1), (i-1,j-1) ]}
</code></pre>

<p>I can definitely take advantage of vectorization to compute the first sum, but I'm having trouble finding a way to vectorize the second sum, since the sum over nearest neighbors depends on the particular set of indices of the lattice site being considered.</p>

<pre><code>def F(a,b, lattice, sites):
    """
    Parameters
    ----------
    a: float
    b: float 
    lattice: np.ndarray, shape [nx, ny]
    sites: dict
         Keys are tuples of the lattice site indices, returns a list 
         of tuples of the site's nearest neighbors

    Returns
    -------
    c: float
    """

    #-- Compute the first sum
    c = a * np.sum( (np.ma.masked_equal(lattice,0) + 1) / 2. )

    #-- Compute the second sum
    c += np.array([ [ b*lattice[i,j]*lattice[x,y] for (x,y) in sites[(i,j)] ] for (i,j) in sites.keys() ]).sum()

    return c
</code></pre>

<p>This gives <em>ok</em> performance, but when I consider large lattices, calculating that second sum with the list comprehension really starts to kill my performance.</p>

<p>Do you guys have any idea how I could go about vectorizing the second sum?</p>

<p>I was thinking a bit about somehow using the strides of the numpy array to calculate this but I'm not if that is (a) feasible and (b) worthwhile.  </p>