построить матрицу из блоков

У меня есть объект, который описывается двумя величинами: A и B (в реальном случае их может быть больше двух). Объекты коррелируются в зависимости от значения A и B. В частности, я знаю корреляционную матрицу для A и для B. Как пример:

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] 

корреляция для A:

корреляция A

поэтому, если элемент имеет A == 0.5, а другой равен A == 1.5, они полностью коррелированы (красные). В противном случае, если элемент имеет A == 0.5, а второй элемент имеет A == 3.5, они некоррелированы (синие).

Аналогично для B:

корреляция B

Теперь я хочу умножить две корреляционные матрицы, но я хочу получить в качестве окончательной матрицы матрицу с двумя осями, где новые оси представляют собой сложенную версию исходных осей:

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

вот что я делаю:

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

визуально:

введите описание изображения здесь

и в частности это должно иметь место:

 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)]) 

на самом деле моя проблема заключается в том, чтобы делать это с большим количеством величин (A, B, C, …) в общем виде. Возможно, есть и более простая функция в numpy для этого.

np.einsum позволяет np.einsum упростить выражение tensordot :

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

Я не думаю, что есть способ устранить reshape .

Также может быть проще обобщить на большее количество массивов, хотя я бы не увлекся слишком большим количеством переменных итерации в одном выражении einsum .

Я думаю, наконец, я нашел решение:

 np.kron(a,b) 

и затем я могу составить

 np.kron(np.kron(a,b), c)