build a numpy array without using for loop

Viewed 117

If I have two numpy array of shape (N, 3) like

array1 = np.array([[x1, y1, 1],
                   ........
                   [xn, yn, 1]])

array2 = np.array([[u1, v1, 1],
                   ........
                   [un, vn, 1]])

I want

A = np.array([[x1 * u1, x1 * v1, x1, y1 * u1, y1 * v1, y1, u1, v1, 1],
              .........
              [xn * un, xn * vn, xn, yn * un, yn * vn, yn, un, vn, 1]]

How should I do it without using for loops to iterate through?

2 Answers

This is called Kronecker product. You can do that using np.kron() function like:

import numpy as np
array1 = np.array([[1,2,1], 
                   [3,4,1]])
array2 = np.array([[3,4,1], 
                   [1,0,1]])
N = array1.shape[0]
result = np.kron(array1, array2)

# A (NxN) x 9 output
#array([[ 3,  4,  1,  6,  8,  2,  3,  4,  1],
#       [ 1,  0,  1,  2,  0,  2,  1,  0,  1],
#       [ 9, 12,  3, 12, 16,  4,  3,  4,  1],
#       [ 3,  0,  3,  4,  0,  4,  1,  0,  1]])

If you want a one to one Kronecker product, you might also use:

# N x 9 
np.kron(array1, array2)[0:N*N:N, :]
#array([[ 3,  4,  1,  6,  8,  2,  3,  4,  1],
#       [ 9, 12,  3, 12, 16,  4,  3,  4,  1]])

which is not recommended if performance matters.

You clearly want a broadcasted multiplication along the second dimension. The trick with all problems like these is to remember that numpy ravels/reshapes in C order by default.

Let's analyze the rows of the output to see how to arrange the product. You have

[x1 * u1, x1 * v1, x1 * 1, y1 * u1, y1 * v1, y1 * 1, u1, v1, 1]

This is the raveled version of

[[x1 * u1, x1 * v1, x1 * 1],
 [y1 * u1, y1 * v1, y1 * 1],
 [ 1 * u1,  1 * v1,  1 * 1]]

Which in turn is the product

array1[0][:, None] * array2[0][None, :]

So the product you want keeps the first dimension and broadcasts the second as shown above. Then you ravel the last two dimensions together to get the partial Kronecker product:

A = (array1[:, :, None] * array2[:, None, :]).reshape(array1.shape[0], -1)
Related