I have an object which is described by two quantities, A and B (in real case they can be more than two). Objects are correlated depending on the value of A and B. In particular I know the correlation matrix for A and for B. Just as example:

a = np.array([[1, 1, 0, 0],[1, 1, 0, 0],[0, 0, 1, 1],[0, 0, 1, 1]])
b = np.array([[1, 1, 0],[1, 1, 1],[0, 1, 1]])
na = a.shape[0]
nb = b.shape[0]

correlation for A:

so if an element has A == 0.5 and the other equal to A == 1.5 they are fully correlated (red). Otherwise if an element has A == 0.5 and the second item has A == 3.5 they are uncorrelated (blue).

Similarly for B:

Now I want multiply the two correlation matrixes, but I want to obtain as final matrix a matrix with two axis, where the new axes are a folded version of the original axes:

def get_folded_bin(ia, ib):return ia * nb + ib

here what I am doing:

result = np.swapaxes(np.tensordot(a, b, axes=0), 1, 2).reshape(na* nb, na * nb)

visually:

and in particular this must hold:

for ia1 in xrange(na):for ia2 in xrange(na):for ib1 in xrange(nb):for ib2 in xrange(nb):assert(a[ia1, ia2] * b[ib1, ib2] == result[get_folded_bin(ia1, ib1), get_folded_bin(ia2, ib2)])

actually my problem is to do it with more quantities (A, B, C, ...) in a general way. Maybe there is also a simpler function within numpy to do that.

np.einsum lets you simplify the tensordot expression a bit:

result = np.einsum('ij,kl->ikjl',a,b).reshape(-1, na * nb)

I don't think there's a way of eliminating the reshape.

It may also be easier to generalize to more arrays, though I wouldn't get carried away with too many iteration variables in one einsum expression.