You are all complicating things unnecessarily... Since you have a single operation, multiplication, which is commutative, you can swap the order in which you perform them at will, i.e. you don't need to multiply items of matrix1
with items of matrix2
, and once you have computed all of them, multiply them together. Instead, you can first multiply all the relevant items of matrix1
together, then all the relevant items of matrix2
together, and then multiply the two resulting values. So you can write your function as the very simple:
def fast_geometric_matrix_multiplication(matrix1, matrix2):
return np.prod(matrix1, axis=1)[:, None] * np.prod(matrix2, axis=0)
It has the additional advantage that, if you are multiplying matrices of shapes (m, k)
and (k, n)
, you would expect to be having to do m*n*2*k
multiplications, while this method only requires m*k + n*k + m*n
, which is bound to be much smaller than what you presently are doing for almost any array shapes.
And of course:
In [24]: a = np.random.rand(100, 200)
...: b = np.random.rand(200, 50)
...:
In [25]: np.allclose(geometric_matrix_multiplication(a, b),
...: fast_geometric_matrix_multiplication(a, b))
Out[25]: True
In [26]: %timeit geometric_matrix_multiplication(a, b)
1 loops, best of 3: 1.39 s per loop
In [27]: %timeit fast_geometric_matrix_multiplication(a, b)
10000 loops, best of 3: 74.2 us per loop
In [28]: %timeit np.prod(a[:, None, :]*b[..., None].T, axis=2)
100 loops, best of 3: 5.97 ms per loop