Vectorized version of array calculation - python-2.7

Question

Is there a way of vectorizing the following array calculation (i.e. without using for loops):
for i in range(numCells):
z[i] = ((i_mask == i)*s_image).sum()/pixel_counts[i]
s_image is an image stored as a 2-dimensional ndarray (I removed the colour dimension here for simplicity). i_mask is also a 2-dimensional array of the same size as s_image but it contains integers which are indexes to a list of 'cells' of length numCells. The result, z, is a 1-dimensional array of length numCells. The purpose of the calculation is to sum all the pixel values where the mask contains the same index and put the results in the z vector. (pixel_counts is also a 1-dimensional array of length numCells).

Answer 1

As one vectorized approach, you can take advantage of broadcasting and matrix-multiplication, like so -
# Generate a binary array of matches for all elements in i_mask against
# an array of indices going from 0 to numCells
matches = i_mask.ravel() == np.arange(numCells)[:,None]
# Do elementwise multiplication against s_image and sum those up for
# each such index going from 0 to numCells. This is essentially doing
# matix multiplicatio. Finally elementwise divide by pixel_counts
out = matches.dot(s_image.ravel())/pixel_counts
Alternatively, as another vectorized approach, you can do those multiplication and summation with np.einsum as well, which might give a boost to the performance, like so -
out = np.einsum('ij,j->i',matches,s_image.ravel())/pixel_counts
Runtime tests -
Function definitions:
def vectorized_app1(s_image,i_mask,pixel_counts):
matches = i_mask.ravel() == np.arange(numCells)[:,None]
return matches.dot(s_image.ravel())/pixel_counts
def vectorized_app2(s_image,i_mask,pixel_counts):
matches = i_mask.ravel() == np.arange(numCells)[:,None]
return np.einsum('ij,j->i',matches,s_image.ravel())/pixel_counts
def org_app(s_image,i_mask,pixel_counts):
z = np.zeros(numCells)
for i in range(numCells):
z[i] = ((i_mask == i)*s_image).sum()/pixel_counts[i]
return z
Timings:
In [7]: # Inputs
...: numCells = 100
...: m,n = 100,100
...: pixel_counts = np.random.rand(numCells)
...: s_image = np.random.rand(m,n)
...: i_mask = np.random.randint(0,numCells,(m,n))
...:
In [8]: %timeit org_app(s_image,i_mask,pixel_counts)
100 loops, best of 3: 8.13 ms per loop
In [9]: %timeit vectorized_app1(s_image,i_mask,pixel_counts)
100 loops, best of 3: 7.76 ms per loop
In [10]: %timeit vectorized_app2(s_image,i_mask,pixel_counts)
100 loops, best of 3: 4.08 ms per loop

Answer 2

Here is my solution (with all three colours handled). Not sure how efficient this is. Anyone got a better solution?
import numpy as np
import pandas as pd
# Unravel the mask matrix into a 1-d array
i = np.ravel(i_mask)
# Unravel the image into 1-d arrays for
# each colour (RGB)
r = np.ravel(s_image[:,:,0])
g = np.ravel(s_image[:,:,1])
b = np.ravel(s_image[:,:,2])
# prepare a dictionary to create the dataframe
data = {'i' : i, 'r' : r, 'g' : g, 'b' : b}
# create a dataframe
df = pd.DataFrame(data)
# Use pandas pivot table to average the colour
# intensities for each cell index value
pixAvgs = pd.pivot_table(df, values=['r', 'g', 'b'], index='i')
pixAvgs.head()
Output:
b g r
i
-1 26.719482 68.041868 101.603297
0 75.432432 170.135135 202.486486
1 92.162162 184.189189 208.270270
2 71.179487 171.897436 201.846154
3 76.026316 178.078947 211.605263

Answer 3

In the end I solved this problem a different way and it drastically increased the speed. Instead of using i_mask as above, a 2-dimensional array of indices into the 1-d array of output intensities, z, I created a different array, mask1593, of dimensions (numCells x 45). Each row is a list of about 35 to 45 indices into the flattened 256x256 pixel image (0 to 65536).
In [10]: mask1593[0]
Out[10]:
array([14853, 14854, 15107, 15108, 15109, 15110, 15111, 15112, 15363,
15364, 15365, 15366, 15367, 15368, 15619, 15620, 15621, 15622,
15623, 15624, 15875, 15876, 15877, 15878, 15879, 15880, 16131,
16132, 16133, 16134, 16135, 16136, 16388, 16389, 16390, 16391,
16392, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)
Then I was able to achieve the same transformation as follows using numpy's advanced indexing:
def convert_image(self, image_array):
"""Convert 256 x 256 RGB image array to 1593 RGB led intensities."""
global mask1593
shape = image_array.shape
img_data = image_array.reshape(shape[0]*shape[1], shape[2])
return np.mean(img_data[mask1593], axis=1)
And here is the result! A 256x256 pixel colour image transformed into an array of 1593 colours for display on this irregular LED display: