Compute weighted sum of 3D and 2D NumPy arrays

Question

Let's say I have a 2D NumPy array A of shape (n, 3) and a 3D array B of shape (x, y, n). I want to create a 3D array Y of shape (x, y, 3) - which is an RGB image.

The 3rd dimension of B contains probabilities [0, 1] which sum up to 1. Array A is a list of RGB colors (one color for each probability). For each element C in the 3rd dimension of B I want to compute the sum of A * C so that I get a RGB color vector.

Example:

Given a colormap A as

A = [(255   0   0),
     (  0 255   0),
     (  0   0 255)]

and an element C in the 3rd dimension of B at point (x, y) as

C = [0.2, 0.6, 0.20]

I want to compute C' = sum(A * C) as

  (255   0   0) * 0.2
+ (  0 255   0) * 0.6
+ (  0   0 255) * 0.2
---------------------
  ( 51 153  51)

and assign

Y[x, y, :] = C'

I know I could just iterate over x and y with a for loop and compute each element at a time but I wonder if this can be done in a vectorized manner that I don't have to iterate over the array myself (mainly for performance reasons).

Answer 1

The method numpy.einsum is convenient for this:

np.einsum('ij,kli->klj', A, B)

The notation says: multiply A[i, j] by B[k, l, i] and sum over i; place the result in the cell [k, l, j].

Example:

A = np.array([(255, 0, 0), (0, 255, 0), (0, 0, 255)])
B = np.array([[[0.2, 0.6, 0.20], [0.2, 0.2, 0.60]], [[0.4, 0.4, 0.2], [0.3, 0.3, 0.4]]])
Y = np.einsum('ij,kli->klj', A, B)

Then Y is

array([[[  51. ,  153. ,   51. ],
        [  51. ,   51. ,  153. ]],

       [[ 102. ,  102. ,   51. ],
        [  76.5,   76.5,  102. ]]])

Answer 2

You are sum-reducing the first axis from A against the third from B, while the rest of the axes are spread out. This is a perfect setup to leverage BLAS based matrix-multiplication for tensors - np.tensordot, like so -

C = np.tensordot(B,A,axes=((2),(0)))

Related post to understand tensordot.

We can also manually reshape to 2D and use the matrix-multiplication for 2D : np.dot, like so -

B.reshape(-1,n).dot(A).reshape(x,y,3)

Note that B.dot(A) works as well, but that would be slower, most probably as it would loop through the first axis of B, while performing 2D matrix-multiplications for each 2D slice off it against A.

Runtime test -

In [180]: np.random.seed(0)
     ...: x,y,n = 100,100,100
     ...: A = np.random.rand(n,3)
     ...: B = np.random.rand(x,y,n)

# @Crazy Ivan's soln
In [181]: %timeit np.einsum('ij,kli->klj', A, B)
100 loops, best of 3: 4.21 ms per loop

In [182]: %timeit np.tensordot(B,A,axes=((2),(0)))
1000 loops, best of 3: 1.72 ms per loop

In [183]: np.random.seed(0)
     ...: x,y,n = 200,200,200
     ...: A = np.random.rand(n,3)
     ...: B = np.random.rand(x,y,n)

# @Crazy Ivan's soln
In [184]: %timeit np.einsum('ij,kli->klj', A, B)
10 loops, best of 3: 33.2 ms per loop

In [185]: %timeit np.tensordot(B,A,axes=((2),(0)))
100 loops, best of 3: 15.3 ms per loop