Numpy/Python: Efficient matrix as multiplication of cartesian product of input matrix

Question

Problem:

The input is a (i,j)-matrix M. The desired output is a (i^n,j^n) matrix K , where n is the number of products taken. The verbose way to get the desired output is as follows

• Generate all arrays of n row permutations I (total of i**n n-arrays)
• Generate all arrays of n column permutations J (total of j**n n-arrays)
• K[i,j] = m[I[0],J[0]] * ... * m[I[n],J[n]] for all n in range(len(J))

The straightforward way I've done this is by generating a list of labels of all n-permutations of numbers in range(len(np.shape(m)[0])) and range(len(np.shape(m)[1])) for rows and columns, respectively. Afterwards you can multiply them as in the last bullet point above. This, however, is not practical for large input matrices -- so I'm looking for ways to optimize the above. Thank you in advance

Example:

Input

np.array([[1,2,3],[4,5,6]])

Output for n = 3

[[ 1. 2. 3. 2. 4. 6. 3. 6. 9. 2. 4. 6. 4. 8. 12. 6. 12. 18. 3. 6. 9. 6. 12. 18. 9. 18. 27.]

[ 4. 5. 6. 8. 10. 12. 12. 15. 18. 8. 10. 12. 16. 20. 24. 24. 30. 36. 12. 15. 18. 24. 30. 36. 36. 45. 54.]

[ 4. 8. 12. 5. 10. 15. 6. 12. 18. 8. 16. 24. 10. 20. 30. 12. 24. 36. 12. 24. 36. 15. 30. 45. 18. 36. 54.]

[ 16. 20. 24. 20. 25. 30. 24. 30. 36. 32. 40. 48. 40. 50. 60. 48. 60. 72. 48. 60. 72. 60. 75. 90. 72. 90. 108.]

[ 4. 8. 12. 8. 16. 24. 12. 24. 36. 5. 10. 15. 10. 20. 30. 15. 30. 45. 6. 12. 18. 12. 24. 36. 18. 36. 54.]

[ 16. 20. 24. 32. 40. 48. 48. 60. 72. 20. 25. 30. 40. 50. 60. 60. 75. 90. 24. 30. 36. 48. 60. 72. 72. 90. 108.]

[ 16. 32. 48. 20. 40. 60. 24. 48. 72. 20. 40. 60. 25. 50. 75. 30. 60. 90. 24. 48. 72. 30. 60. 90. 36. 72. 108.]

[ 64. 80. 96. 80. 100. 120. 96. 120. 144. 80. 100. 120. 100. 125. 150. 120. 150. 180. 96. 120. 144. 120. 150. 180. 144. 180. 216.]]

Partial solution:

The best I've found is a function to create the cartesian product of matrices proposed here: https://stackoverflow.com/a/1235363/4003747 The problem is that the output is not a matrix but an array of arrays. Multiplying the element of each array gives the values I'm after, but in an unordered fashion. I've tried for a while but I have no idea how to sensibly reorder them.

Inefficient solution for n =3:

import numpy as np
import itertools

m=np.array([[1,2,3],[4,5,6]])

def f(m):
    labels_i = [list(p) for p in itertools.product(range(np.shape(m)[0]),repeat=3)]
    labels_j = [list(p) for p in itertools.product(range(np.shape(m)[1]),repeat=3)]


    out = np.zeros([len(labels_i),len(labels_j)])
    for i in range(len(labels_i)):
        for j in range(len(labels_j)):
            out[i,j] = m[labels_i[i][0],labels_j[j][0]] * m[labels_i[i][1],labels_j[j][1]] * m[labels_i[i][2],labels_j[j][2]]
    return out

Answer 1

Here's a vectorized approach using a combination of broadcasting and linear indexing -

from itertools import product

# Get input array's shape
r,c = A.shape

# Setup arrays corresponding to labels i and j
arr_i = np.array(list(product(range(r), repeat=n)))
arr_j = np.array(list(product(range(c), repeat=n)))

# Use linear indexing with ".ravel()" to extract elements.
# Perform elementwise product along the rows for the final output
out = A.ravel()[(arr_i*c)[:,None,:] + arr_j].prod(2)

Runtime test and output verification -

In [167]: # Inputs          
     ...: n = 4
     ...: A = np.array([[1,2,3],[4,5,6]])
     ...: 
     ...: def f(m):
     ...:   labels_i = [list(p) for p in product(range(np.shape(m)[0]),repeat=n)]
     ...:   labels_j = [list(p) for p in product(range(np.shape(m)[1]),repeat=n)]
     ...: 
     ...:   out = np.zeros([len(labels_i),len(labels_j)])
     ...:   for i in range(len(labels_i)):
     ...:      for j in range(len(labels_j)):
     ...:          out[i,j] = m[labels_i[i][0],labels_j[j][0]] \
     ...:                   * m[labels_i[i][1],labels_j[j][1]] \
     ...:                   * m[labels_i[i][2],labels_j[j][2]] \
     ...:                   * m[labels_i[i][3],labels_j[j][3]]
     ...:   return out
     ...: 
     ...: def f_vectorized(A,n):
     ...:   r,c = A.shape
     ...:   arr_i = np.array(list(product(range(r), repeat=n)))
     ...:   arr_j = np.array(list(product(range(c), repeat=n)))
     ...:   return A.ravel()[(arr_i*c)[:,None,:] + arr_j].prod(2)
     ...: 

In [168]: np.allclose(f_vectorized(A,n),f(A))
Out[168]: True

In [169]: %timeit f(A)
100 loops, best of 3: 2.37 ms per loop

In [170]: %timeit f_vectorized(A,n)
1000 loops, best of 3: 202 µs per loop