Let's experiment with 2 small arrays:
In [124]: A, B = np.array([[1,2],[3,4]]), np.array([[10,11],[12,13]])
kron
produces:
In [125]: np.kron(A,B)
Out[125]:
array([[10, 11, 20, 22],
[12, 13, 24, 26],
[30, 33, 40, 44],
[36, 39, 48, 52]])
outer
produces the same numbers, but with a different arangement:
In [126]: np.outer(A,B)
Out[126]:
array([[10, 11, 12, 13],
[20, 22, 24, 26],
[30, 33, 36, 39],
[40, 44, 48, 52]])
kron
reshapes it to a combination of the shapes of A
and B
:
In [127]: np.outer(A,B).reshape(2,2,2,2)
Out[127]:
array([[[[10, 11],
[12, 13]],
[[20, 22],
[24, 26]]],
[[[30, 33],
[36, 39]],
[[40, 44],
[48, 52]]]])
it then recombines 4 dimensions into 2 with concatenate
:
In [128]: np.concatenate(np.concatenate(_127, 1),1)
Out[128]:
array([[10, 11, 20, 22],
[12, 13, 24, 26],
[30, 33, 40, 44],
[36, 39, 48, 52]])
An alternative is to swap axes, and reshape:
In [129]: _127.transpose(0,2,1,3).reshape(4,4)
Out[129]:
array([[10, 11, 20, 22],
[12, 13, 24, 26],
[30, 33, 40, 44],
[36, 39, 48, 52]])
The first reshape and transpose produce a view, but the second reshape has to produce a copy. Concatenate makes a copy. But all those actions are done in compiled numpy
code.
Defining functions:
def foo1(A,B):
temp = np.outer(A,B)
temp = temp.reshape(A.shape + B.shape)
return np.concatenate(np.concatenate(temp, 1), 1)
def foo2(A,B):
temp = np.outer(A,B)
nz = temp.shape
temp = temp.reshape(A.shape + B.shape)
return temp.transpose(0,2,1,3).reshape(nz)
testing:
In [141]: np.allclose(np.kron(A,B), foo1(A,B))
Out[141]: True
In [142]: np.allclose(np.kron(A,B), foo2(A,B))
Out[142]: True
timing:
In [143]: timeit np.kron(A,B)
42.4 μs ± 294 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [145]: timeit foo1(A,B)
26.3 μs ± 38.6 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [146]: timeit foo2(A,B)
13.8 μs ± 19.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
My code may need some generalization, but it demonstrates the validity of the approach.
===
With your kron
:
In [150]: kron(A,B)
Out[150]:
array([[10., 11., 20., 22.],
[12., 13., 24., 26.],
[30., 33., 40., 44.],
[36., 39., 48., 52.]])
In [151]: timeit kron(A,B)
55.3 μs ± 1.59 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
edit
einsum
can do both the outer and transpose:
In [265]: np.einsum('ij,kl->ikjl',A,B).reshape(4,4)
Out[265]:
array([[10, 11, 20, 22],
[12, 13, 24, 26],
[30, 33, 40, 44],
[36, 39, 48, 52]])
In [266]: timeit np.einsum('ij,kl->ikjl',A,B).reshape(4,4)
9.87 μs ± 33 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)