Доступно ли многоуровневое вещание в NumPy?

Question

Интересно, есть ли встроенная операция, которая освободила бы мой код от циклов Python.

Проблема заключается в следующем: у меня есть две матрицы A а также B, A имеет N строки и B имеет N колонны. Я хотел бы умножить каждый i ряд из A с соответствующими i колонка из B (используя трансляцию NumPy). Результирующая матрица будет формироваться i слой на выходе. Так что мой результат будет 3-мерный массив.

Доступна ли такая операция в NumPy?

Answer 1

Да, в простейшей форме вы просто добавляете "нулевые" измерения, чтобы NumPy транслировал по строкам A и столбцы B:

>>> import numpy as np

>>> A = np.arange(12).reshape(3, 4) # 3 row, 4 colums
>>> B = np.arange(15).reshape(5, 3) # 5 rows, 3 columns
>>> res = A[None, ...] * B[..., None]
>>> res
array([[[  0,   0,   0,   0],
        [  4,   5,   6,   7],
        [ 16,  18,  20,  22]],

       [[  0,   3,   6,   9],
        [ 16,  20,  24,  28],
        [ 40,  45,  50,  55]],

       [[  0,   6,  12,  18],
        [ 28,  35,  42,  49],
        [ 64,  72,  80,  88]],

       [[  0,   9,  18,  27],
        [ 40,  50,  60,  70],
        [ 88,  99, 110, 121]],

       [[  0,  12,  24,  36],
        [ 52,  65,  78,  91],
        [112, 126, 140, 154]]])

Результат имеет форму (5, 3, 4) и вы можете легко перемещать ось, если хотите другую форму. Например, используя np.moveaxis:

>>> np.moveaxis(res, (0, 1, 2), (2, 0, 1))  # 0 -> 2 ; 1 -> 0, 2 -> 1
array([[[  0,   0,   0,   0,   0],
        [  0,   3,   6,   9,  12],
        [  0,   6,  12,  18,  24],
        [  0,   9,  18,  27,  36]],

       [[  4,  16,  28,  40,  52],
        [  5,  20,  35,  50,  65],
        [  6,  24,  42,  60,  78],
        [  7,  28,  49,  70,  91]],

       [[ 16,  40,  64,  88, 112],
        [ 18,  45,  72,  99, 126],
        [ 20,  50,  80, 110, 140],
        [ 22,  55,  88, 121, 154]]])

С формой (3, 4, 5),

Answer 2

Один из способов выразить ваше требование напрямую - использовать np.einsum():

>>> A = np.arange(12).reshape(3, 4)
>>> B = np.arange(15).reshape(5, 3)
>>> np.einsum('...i,j...->...ij', A, B)
array([[[  0,   0,   0,   0,   0],
        [  0,   3,   6,   9,  12],
        [  0,   6,  12,  18,  24],
        [  0,   9,  18,  27,  36]],

       [[  4,  16,  28,  40,  52],
        [  5,  20,  35,  50,  65],
        [  6,  24,  42,  60,  78],
        [  7,  28,  49,  70,  91]],

       [[ 16,  40,  64,  88, 112],
        [ 18,  45,  72,  99, 126],
        [ 20,  50,  80, 110, 140],
        [ 22,  55,  88, 121, 154]]])

Это использует соглашение суммирования Эйнштейна.

Дальнейшее обсуждение см. В главе 3 " Векторы, чистый и прикладной: общее введение в линейную алгебру " Т. В. Кёрнера. В нем автор приводит забавный отрывок из письма Эйнштейна другу:

"Я сделал великое открытие в математике; я подавлял знак суммирования каждый раз, когда суммирование должно производиться по индексу, который встречается дважды..."