Is it possible to vectorize a function that access different elements in an numpy array?

Question

Let's say that I wish to implement the following code in python

This function takes an image as 1 dimensional array, and iterates over individual elements in the array (pixels in the input image) which affects an output array that is also an image represented as a 1 dimensional array

example: a single pixel in the input image (red) affects 8 surrounding pixels in (orange)

A basic implementation in C is

/* C version 
 * Given an input image create an 
 * output image that is shaped by individual pixels
 * of the input image
 */

int image[width * height]; //image retrieved elsewhere
int output [width * height]; //output image
int y = 0, x = 0;
for( y = 1; y < height-1 ; ++ y) {
    for(x = 1; x < width-1; ++ x) {
        if (image[y * width + x] > condition) {
            /* pixel affects the surrounding 8 pixels in the output image */

            output[(y-1) * width + x - 1]++; /* upper left  */
            output[(y-1) * width + x    ]++; /* above       */
            output[(y-1) * width + x + 1]++; /* upper right */
            output[y * width + x + 1    ]++; /* right       */
            output[y * width + x - 1    ]++; /* left        */
            output[(y+1) * width + x - 1]++; /* lower left  */
            output[(y+1) * width + x    ]++; /* below       */
            output[(y+1) * width + x + 1]++; /* lower right */


        }
    }
}

the naive approach in python would be to use the exactly the same element wise access as shown below

#Python version
input  = blah # formed elsewhere  
output = np.zeros(width * height)
for y in xrange(1, height-1):
    for x in xrange(1, width-1):
        if input[y * width + x] > condition:
            output[(y-1) * width + x - 1]+= 1; # upper left  
            output[(y-1) * width + x    ]+= 1; # above       
            output[(y-1) * width + x + 1]+= 1; # upper right 
            output[y * width + x + 1    ]+= 1; # right       
            output[y * width + x - 1    ]+= 1; # left        
            output[(y+1) * width + x - 1]+= 1; # lower left  
            output[(y+1) * width + x    ]+= 1; # below       
            output[(y+1) * width + x + 1]+= 1; # lower right 

is there a better way to implement this? Is is possible to vectorize this function?

Answer 1

If I understood the question correctly, then the approach can be turned upside down: if a pixel has in its neighbourhood pixels that match the condition, increment it by one for each match. Do this for all pixels. Scipy (among others) offers tools for filtering images:

In [51]: import scipy.ndimage

Create the sample image from the 1-dimensional array. Reshape creates a view instead of copying:

In [62]: I1d
Out[62]: 
array([  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
         0,   0,   0,   0,   0, 129,   0, 129, 129,   0,   0,   0,   0,
         0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 129,   0,   0,
         0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
         0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 129])

In [63]: height
Out[63]: 8

In [64]: width
Out[64]: 8

In [65]: I = I1d.reshape((height, width))

In [66]: I
Out[66]: 
array([[  0,   0,   0,   0,   0,   0,   0,   0],
       [  0,   0,   0,   0,   0,   0,   0,   0],
       [  0,   0, 129,   0, 129, 129,   0,   0],
       [  0,   0,   0,   0,   0,   0,   0,   0],
       [  0,   0,   0,   0, 129,   0,   0,   0],
       [  0,   0,   0,   0,   0,   0,   0,   0],
       [  0,   0,   0,   0,   0,   0,   0,   0],
       [  0,   0,   0,   0,   0,   0,   0, 129]])

Using convolution create an image that holds the increments to each pixel in the original from a binary mask of pixels that exceed the condition (128 here):

In [67]: scipy.ndimage.filters.convolve(
    (I > 128).astype(np.int),  # conditioned binary image
    weights=np.array([[1, 1, 1],  # each match weighted as 1
                      [1, 0, 1],
                      [1, 1, 1]]),
    mode='constant', cval=0)  # Use zeros as constant fill values on edges
Out[67]: 
array([[0, 0, 0, 0, 0, 0, 0, 0],
       [0, 1, 1, 2, 2, 2, 1, 0],
       [0, 1, 0, 2, 1, 1, 1, 0],
       [0, 1, 1, 3, 3, 3, 1, 0],
       [0, 0, 0, 1, 0, 1, 0, 0],
       [0, 0, 0, 1, 1, 1, 0, 0],
       [0, 0, 0, 0, 0, 0, 1, 1],
       [0, 0, 0, 0, 0, 0, 1, 0]])

In [68]: conv = _

If the final goal is to add the original and the increments:

In [69]: I + conv
Out[69]: 
array([[  0,   0,   0,   0,   0,   0,   0,   0],
       [  0,   1,   1,   2,   2,   2,   1,   0],
       [  0,   1, 129,   2, 130, 130,   1,   0],
       [  0,   1,   1,   3,   3,   3,   1,   0],
       [  0,   0,   0,   1, 129,   1,   0,   0],
       [  0,   0,   0,   1,   1,   1,   0,   0],
       [  0,   0,   0,   0,   0,   0,   1,   1],
       [  0,   0,   0,   0,   0,   0,   1, 129]])

To output a 1-dimensional array as well, use either ravel() or flatten(). The former creates a 1-dimensional view to the original 2-dimensional array, the latter creates a flattened copy:

In [70]: conv.ravel()
Out[70]: 
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 1, 0, 0, 1, 0, 2, 1, 1, 1,
       0, 0, 1, 1, 3, 3, 3, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1,
       0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0])

Answer 2

I think what you are trying to do is easiest with 2D array indexing. You can reshape arrays easily with numpy. No data is copied. The new 2D array just provides a convenient way to index the same values that are stored in the original array. Here is an example code.

#imports
import numpy as np


# Setup
Nx = 5
Ny = 7
cutoff = 3.0
arr_input = np.array([[0, 0, 0, 0, 0, 0, 9],
                      [0, 4, 0, 0, 0, 0, 0], 
                      [0, 0, 0, 0, 0, 3, 0], 
                      [0, 2, 0, 0, 0, 0, 1], 
                      [0, 0, 0, 0, 5, 0, 0]]).flatten()

# Make an array of center locations with input value bigger than the cutoff value 
centers_array_2d = np.where(arr_input>=cutoff, 1.0, 0.0)

# Initialize the output array
arr_output = np.zeros_like(centers_array_2d)

# Reshape the arrays to use 2D indexing
ai = centers_array_2d.reshape(Nx, Ny)
ao = arr_output.reshape(Nx, Ny)

# Do the neighbor calculation with numpy indexing rules
ao[:-1, :-1] += ai[1:, 1:]               # lower left
ao[:, :-1] += ai[:, 1:]                  # lower center
ao[1:, :-1] += ai[:-1, 1:]               # lower right
ao[:-1, :] += ai[1:, :]                  # middle left
# ao[:, :] += ai[:, :]                   # middle center 
ao[1:, :] += ai[:-1, :]                  # middle right
ao[:-1, 1:] += ai[1:, :-1]               # top left
ao[:, 1:] += ai[:, :-1]                  # top center
ao[1:, 1:] += ai[:-1, :-1]               # top right

# Reshape the output array to return a 1D array like the input
arr_output = ao.flatten()

# Print the results 
print('input1d: \n{}\n'.format(arr_input))
print("2D array 'ai':\n{}\n".format(ai))
print("2D array 'ao':\n{}\n".format(ao))
print('output1d: \n{}\n'.format(arr_output))

The arrays are as follows.

input1d: 
[0 0 0 0 0 0 9 0 4 0 0 0 0 0 0 0 0 0 0 3 0 0 2 0 0 0 0 1 0 0 0 0 5 0 0]

2D array 'ai':
[[ 0.  0.  0.  0.  0.  0.  1.]
 [ 0.  1.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.  1.  0.]
 [ 0.  0.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  1.  0.  0.]]

2D array 'ao':
[[ 1.  1.  1.  0.  0.  1.  0.]
 [ 1.  0.  1.  0.  1.  2.  2.]
 [ 1.  1.  1.  0.  1.  0.  1.]
 [ 0.  0.  0.  1.  2.  2.  1.]
 [ 0.  0.  0.  1.  0.  1.  0.]]

output1d: 
[ 1.  1.  1.  0.  0.  1.  0.  1.  0.  1.  0.  1.  2.  2.  1.  1.  1.  0. 1.  0.  1.  0.  0.  0.  1.  2.  2.  1.  0.  0.  0.  1.  0.  1.  0.]

Is this the calculation that you were looking for? This is what I would interpret as vectorizing the code you gave. Also, you could make a list of indices for the 1D arrays that correspond to each neighbor. This is essentially what is going on internally when I use 2D slice indices to access elements of the 2D arrays.