Is it possible to access the current indices during a Numpy vectorized broadcasting operation?

Question

I would like to speed up a function on a single array in Numpy using fancy indexing, vectorization, and/or broadcasting. For each value in my array, I need to do a calculation that involves adjacent values. Therefore, in my vectorized operation, I need to have access to the current index so that I can grab indices around it. Consider the following simple array operation:

x = np.arange(36).reshape(6, 6)
y = np.zeros((6, 6))
y[:] = x + 1

I'd like to use similar syntax, but rather than a simple increment, I'd like to do something like add all values at adjacent indices to the current value in the vectorized loop. For instance if the region around index [i, j] == 7 looks like

3 2 5
2 7 6
5 5 5

I'd like the calculated value for [i, j] to be 3 + 2 + 5 + 2 + 7 + 6 + 5 + 5 + 5 , and I want to do that for all indices [i, j] .

This is a straightforward nested for loop (or a single for loop using np.sum for each index)... but I want to use broadcasting and/or fancy indexing if possible. This may be too complex of a problem for the Numpy syntax, but I feel like it should be possible.

Essentially, it comes down to this: how do I reference the current index during a broadcasting operation?

Answer 1

Start with a 1D example:

x = np.arange(10)

There is a choice you have to make: do you discard the edges or not, since they don't have two neighbors? If you do, you can create your output array in esentially one step:

result = x[:-2] + x[1:-1] + x[2:]

Notice that all three addends are views because they use simple indexing. You want to avoid fancy indexing as much as you can because it generally involves making copies.

If you prefer to retain the edges, you can pre-allocate the output buffer and add directly into it:

result = x.copy()
result[:-1] += x[1:]
result[1:] += x[:-1]

The fundamental idea in both cases is that to apply an operation to all neighboring elements, you just shift the array by +/-1. You don't need to know any indices, or do anything fancy. The simpler the better.

Hopefully you can see how how to generalize this to the 2D case. Rather than a single index shifting between -1, 0, 1, you have two indices in every possible combination of -1, 0, 1 between the two of them.

Appendix

Here's the generalized approach for a no-egde result:

from itertools import product
def sum_shifted(a):
    result = np.zeros(tuple(x - 2 for x in a.shape), dtype=a.dtype)
    for index in product([slice(0, -2), slice(1, -1), slice(2, None)], repeat=a.ndim):
        result += a[index]
    return result

This implementation is somewhat rudimentary because it doesn't check for inputs with no dimensions or shapes < 2, but it does work for arbitrary numbers of dimensions.

Notice that for a 1D case, the loop will run exactly three times, for 2D nine times and for ND 3N. This is one case where I find an explicit for loop to be appropriate with numpy. The loop is very small compared to the work done on a large array, fast enough for a small array, and certainly better than writing all 27 possibilities out by hand for the 3D case.

One more thing to pay attention to is how the successive indices are generated. In Python an index with a colon, like x[1:2:3] is converted to the relatively unknown slice object: slice(1, 2, 3) . Since (almost) everything with commas gets interpreted as a tuple, an index like in the expression x[1:2, ::-1, :2] is exactly equivalent to (slice(1, 2), slice(None, None, -1), slice(None, 2)) . The loop generates exactly such an expression, with one element for each dimension. So the result is actually simple indexing across all dimensions.

A similar approach is possible if you want to retain edges. The only significant difference is that you need to index both the input and the output arrays:

from itertools import product
def sum_shifted(a):
    result = np.zeros_like(a)
    for r_index, a_index in zip(product([slice(0, -1), slice(None), slice(1, None)], repeat=a.ndim),
                                product([slice(1, None), slice(None), slice(0, -1)], repeat=a.ndim)):
        result[r_index] += a[a_index]
    return result

This works because itertools.product guarantees the order of the iteration, so the two zipped iterators will stay in lockstep.

Answer 2

try this:

x = np.arange(36).reshape(6, 6)
y = np.zeros((6, 6))
for i in range(x.shape[0]):
    for j in range(x.shape[1]):
        if i>0 and i<x.shape[0]-1 and j>0 and j<x.shape[1]-1:
            y[i,j]=x[i,j]+x[i-1,j]+x[i,j-1]+x[i-1,j-1]+x[i+1,j]+x[i,j+1]+x[i+1,j+1]+x[i-1,j+1]+x[i+1,j-1]
        if j==0:
            if i==0:
                y[i,j]=x[i,j]+x[i,j+1]+x[i+1,j+1]+x[i+1,j]
            elif i==x.shape[0]-1:
                y[i,j]=x[i,j]+x[i,j+1]+x[i-1,j+1]+x[i-1,j]
            else:
                y[i,j]=x[i,j]+x[i,j+1]+x[i+1,j+1]+x[i+1,j]+x[i-1,j]+x[i-1,j+1]

        if j==x.shape[1]-1:
            if i==0:
                y[i,j]=x[i,j]+x[i,j-1]+x[i+1,j-1]+x[i+1,j]
            elif i==x.shape[0]-1:
                y[i,j]=x[i,j]+x[i,j-1]+x[i-1,j-1]+x[i-1,j] 
            else:
                y[i,j]=x[i,j]+x[i,j-1]+x[i-1,j-1]+x[i+1,j]+x[i-1,j]+x[i+1,j-1]
        if i==0 and j in range(1,x.shape[1]-1):
            y[i,j]=x[i,j]+x[i,j-1]+x[i+1,j-1]+x[i+1,j]+x[i+1,j+1]+x[i,j+1]
        if i==x.shape[0]-1 and j in range(1,x.shape[1]-1):
            y[i,j]=x[i,j]+x[i,j-1]+x[i-1,j-1]+x[i-1,j]+x[i-1,j+1]+x[i,j+1]
print(y)