11. Транслирование массивов

Иногда бывает необходимо выполнять арифметические операции c массивами разной формы. Допустим, мы хотим умножить один массив на другой, причем размерность у массивов будет разной:

>>> a = np.arange(9).reshape(3, 3)
>>> a
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>>
>>> b = np.array([3, 5, 7])
>>> b
array([3, 5, 7])
>>>
>>> c = a*b
>>> c
array([[ 0,  5, 14],
       [ 9, 20, 35],
       [18, 35, 56]])

С точки зрения линейной алгебры - это бред. Но на самом деле в таком умножении есть довольно много логики: каждый столбец из первого массива был умножен на соответствующий элемент из второго массива. Если бы подобный механизм в NumPy отсутствовал, нам пришлось бы делать подобные операции вручную, каждый раз изобретая что-то подобное:

import numpy as np

a = np.arange(9).reshape(3, 3)
b = np.array([3, 5, 7])

c = np.empty(a.shape)

for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        c[i][j] = a[i][j]*b[j]

И дело не в том, что нам пришлось бы тратить больше времени, а в том что нам пришлось бы это делать на Python, а скоростью этот язык особо похвастаться не может. Но благодаря транслированию весь этот механизм работает под капотом, причем на языке С, что на порядки быстрее.

Транслирование для двух (и более) массивов выполняется по двум правилам:

  • если массивы имеют различное количество осей (измерений), то к размерностям массивов с недостающими осями будет добавляться единица, до тех пор пока размерности (количество осей) двух массивов не совпадут;
  • если по какой-то из осей находится всего один элемент, то эта ось будет вести себя так, как если бы в ней был не один элемент, а ровно столько, сколько элементов в соответствующей оси другого массива.

Но прежде чем говорить о применении этих правил в транслировании давайте сначала разберем несколько простых примеров. Умножим одномерный массив на число (скаляр).

>>> a = np.arange(1, 11)
>>> a
array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])
>>>
>>> b = 2
>>> 
>>> c = b*a
>>> c
array([ 2,  4,  6,  8, 10, 12, 14, 16, 18, 20])

Каждый элемент массива a умножился на число b - ничего сложного. Но давайте перепишем этот пример немного иначе.

>>> a = np.arange(1, 11)
>>> b = np.array([2, 2, 2, 2, 2, 2, 2, 2, 2, 2])
>>>
>>> c = b*a
>>> c
array([ 2,  4,  6,  8, 10, 12, 14, 16, 18, 20])

Результат не изменился. Но во втором случае мы использовали гораздо больше памяти, что не очень-то эффективно. Теперь давайте попробуем что-нибудь посложнее - умножим каждый столбец двумерного массива на соответствующий элемент из другого, одномерного массива.

>>> a = np.arange(9).reshape(3, 3)
>>> b = np.array([3, 5, 7])
>>>
>>> a
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> b
array([3, 5, 7])
>>> 
>>> #  Создадим массив из нулей с формой массива "а"
...
>>> c = np.empty(a.shape)
>>> 
>>> for i in range(a.shape[0]):
...     for j in range(a.shape[1]):
...         c[i][j] = a[i][j]*b[j]
... 
>>> c
array([[  0.,   5.,  14.],
       [  9.,  20.,  35.],
       [ 18.,  35.,  56.]])

Но благодаря транслированию NumPy позволяет нам записать, просто, вот так:

>>> a = np.arange(9).reshape(3, 3)
>>> b = np.array([3, 5, 7])
>>>
>>> a
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> b
array([3, 5, 7])
>>>
>>> c = b*a
>>> c
array([[ 0,  5, 14],
       [ 9, 20, 35],
       [18, 35, 56]])

Даже интуитивно понятно, что второй вариант гораздо лучше первого. И это действительно так: во втором варианте нам пришлось гораздо меньше стучать по клавиатуре чем в первом варианте; в первом варианте мы пишем алгоритм умножения двух массивов на Python в то время как во втором варианте умножение происходит под капотом NumPy на языке С, что намного быстрее. А теперь представьте, что у нас есть трехмерный массив a

>>> a = np.arange(12).reshape(2, 2, 3)
>>> a
array([[[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 6,  7,  8],
        [ 9, 10, 11]]])

И к каждому его элементу нам нужно прибавить единицу. Я думаю, что теперь понятно насколько это просто сделать в NumPy:

>>> c = a + 1
>>> c
array([[[ 1,  2,  3],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [10, 11, 12]]])

Не все приведенные примеры имеют отношение к механизму транслирования массивов, но они хорошо отражают его суть: благодаря ему мы можем выполнять арифметические операции с массивами разного размера.


11.1. Правила транслирования

Итак, благодаря механизму транслирования мы можем выполнять вполне логичные математические операции над массивами разного размера. Причем, это выполняется внутренними алгоритмами реализованными на языке С, а это гораздо быстрее чем на Питоне. Возникает вопрос: когда транслирование работает, а когда нет?

В NumPy существует понятие - совместимость размеров массивов. Два размера считаются совместимыми если они равны или один из размеров равен 1. Например:

>>> a = np.arange(4).reshape(2, 2)
>>> a
array([[0, 1],
       [2, 3]])
>>> b = np.array([10])
>>>
>>> a + b
array([[10, 11],
       [12, 13]])

В данном случае массив b транслируется по массиву a, т.е. размеры b приводятся в соответствие к размерам массива a. В самом деле:

>>> b.ndim    #  У массива всего одна ось
1
>>> b.shape    #  и всего один элемент
(1,)
>>>
>>> b = np.hstack((b, np.array([10])))     # Расширяем "b" по горизонтали
>>> b    #  и теперь размеры массивов "а" и "b" по горизонтали совпадают
array([10, 10])
>>>
>>> a + b    #  Это снова сработает
array([[10, 11],
       [12, 13]])
>>>
>>> b = np.vstack((b, np.array([10, 10])))     #  Потому что "b" может быть  
>>> b                                          #  увеличен по вертикали
array([[10, 10],
       [10, 10]])
>>>
>>> b.ndim    #  Теперь количество осей
2
>>> b.shape    # и размеры вдоль этих осей у массивов совпадают
(2, 2)

Приблизительно так и работает правило совместимости размеров массивов. В примере выше размеры и количество осей не совпадали, но размер одного массива оказался равен 1, что и позволило произвести его расширение до нужных размеров. Но давайте разберем пример посложнее.

>>> a = np.ones(4).reshape(2, 2)
>>> a
array([[ 1.,  1.],
       [ 1.,  1.]])
>>>
>>> b = np.arange(6).reshape(3, 1, 2)
>>> b
array([[[0, 1]],

       [[2, 3]],

       [[4, 5]]])
>>>
>>> c = b + a
>>> c    #  Работает!
array([[[ 1.,  2.],
        [ 1.,  2.]],

       [[ 3.,  4.],
        [ 3.,  4.]],

       [[ 5.,  6.],
        [ 5.,  6.]]])
>>> 
>>> a.shape
(2, 2)
>>> b.shape
(3, 1, 2)
>>> c.shape
(3, 2, 2)
>>>
>>> #  Последние оси совпадают по размерам
... #  b: 3 x 1 x 2
... #  a:     2 x 2
... #  c: 3 x 2 x 2
... #    _----^----_
... #  предпоследняя ось "b" может быть расширена до 2
... 
>>> b = np.arange(6).reshape(3, 2, 1)    #  Изменим размеры "b"
>>> b
array([[[0],
        [1]],

       [[2],
        [3]],

       [[4],
        [5]]])
>>>
>>> c = b + a    #  Снова работает!
>>> c
array([[[ 1.,  1.],
        [ 2.,  2.]],

       [[ 3.,  3.],
        [ 4.,  4.]],

       [[ 5.,  5.],
        [ 6.,  6.]]])
>>> 
>>> a.shape
(2, 2)
>>> b.shape
(3, 2, 1)
>>> c.shape
(3, 2, 2)
>>>
>>> #  b: 3 x 2 x 1
... #  a:     2 x 2
... #  c: 3 x 2 x 2
... #        _----^----_
... #  последняя ось "b" может быть расширена до 2
... 
>>> b = np.arange(6).reshape(2, 3, 1)    #  Снова изменим размер "b"
>>> b
array([[[0],
        [1],
        [2]],

       [[3],
        [4],
        [5]]])
>>> c = b + a
Traceback (most recent call last):
  File "", line 1, in 
ValueError: operands could not be broadcast 
together with shapes (2,3,1) (2,2) 
>>> 
>>> #  b: 2 x 3 x 1
... #  a:     2 x 2
... #    _----^----_
... #  эти размеры не совпадают

Можно сказать: "Ну и что, что не совпадают. Массив а состоит из одних единиц, что нам мешает добавить еще одну строку?". Но то, что массив а состоит из одних единиц (одинаковых элементов) - случайность. Такое расширение массивов с произвольными элементами привело бы к появлению фиктивных данных.

>>> #  Как бы добавим еще одну строку к массиву "а"
...
>>> a = np.ones(6).reshape(3, 2)
>>> b = np.arange(6).reshape(2, 3, 1)
>>>
>>> a
array([[ 1.,  1.],
       [ 1.,  1.],
       [ 1.,  1.]])
>>> b
array([[[0],
        [1],
        [2]],

       [[3],
        [4],
        [5]]])
>>>
>>> c = a + b    #  ... и все работает
>>> c
array([[[ 1.,  1.],
        [ 2.,  2.],
        [ 3.,  3.]],

       [[ 4.,  4.],
        [ 5.,  5.],
        [ 6.,  6.]]])
>>>
>>> a.shape
(3, 2)
>>> b.shape
(2, 3, 1)
>>> c.shape
(2, 3, 2)
>>>
>>> #  b: 2 x 3 x 1
... #  a:     3 x 2
... #  c: 2 x 3 x 2
... #        _----^----_
... #  последняя ось массива "b" может быть расширена до 2

Мы можем поступать таким образом с массивами, которые состоят из одинаковых элементов, т.е. добавить необходимое количество строк или столбцов из тех же самых элементов. Но если массив состоит из разных элементов, то добавить строки и столбцы уже не получится. Например:

>>> a = np.array([[5, 5],[5, 5]])
>>>
>>> #  К такому массиву легко добавить 
... #  строку или столбец из одних пятерок
...
>>> a
array([[5, 5],
       [5, 5]])
>>>
>>> a = np.array([[1, 18],[3, 42]])
>>>
>>> #  Какие тут можно выдумать строки или
... # столбцы и добавить их - непонятно
>>> a
array([[ 1, 18],
       [ 3, 42]])

И так, в случае, когда количество осей совпадает, и некоторые из осей содержат всего по одному элементу, а количество элементов по другим осям одинаково, то такие масивы имеют совместимые размеры. А раз они имеют совместимые размеры, значит может быть выполнено второе правило транслирования: оси на которых всего один элемент ведут себя так, буд-то там не один элемент, а столько сколько на соответствующей оси другого массива.

>>> a = np.arange(48).reshape(2, 3, 1, 8)
>>> b = np.arange(96).reshape(2, 3, 2, 8)
>>> c = a + b
>>> c.shape
(2, 3, 2, 8)
>>>
>>> #  a: 2 x 3 x 1 x 8
... #  b: 2 x 3 x 2 x 8
... #  c: 2 x 3 x 2 x 8
... #     --------^----
... #  третья ось "a" может быть расширена до 2
... 
>>> a = np.arange(48).reshape(2, 3, 1, 8)
>>> b = np.arange(18).reshape(2, 1, 9, 1)
>>> c = a + b
>>> c.shape
(2, 3, 9, 8)
>>>
>>> #  a: 2 x 3 x 1 x 8
... #  b: 2 x 1 x 9 x 1
... #  c: 2 x 3 x 9 x 8
... #     ----^---^---^
... #  2-я ось расширена до 3, 3-я ось до 9 и 4-я до 8
... 

Но всегда ли количество осей массивов можно привести к одинаковому количеству? Ответ - нет.

>>> a = np.arange(48).reshape(2, 3, 1, 8)
>>> b = np.arange(2).reshape(2, 1)
>>>
>>> c = a + b
>>> c.shape
(2, 3, 2, 8)
>>>
>>> #  a: 2 x 3 x 1 x 8
... #  b:(1)x(1)x 2 x 1    #  NumPy транслирует "b", как 1 x 1 x 2 x 1 где
... #  c: 2 x 3 x 2 x 8    #  первая, вторая и четвертая оси могут быть 
... #     --------^---^    #  расширены до соответствующих размеров массива
...                        #  "а". Третья ось "а" - расширяется до 2.
>>> #  в самом деле:
... 
>>> a = np.arange(48).reshape(2, 3, 1, 8)
>>> b = np.arange(2).reshape(1, 1, 2, 1)
>>>
>>> c = a + b
>>> c.shape
(2, 3, 2, 8) 
>>>
>>> #  Давайте попробуем такой вариант:
... 
>>> a = np.arange(48).reshape(2, 3, 1, 8)
>>> b = np.arange(3).reshape(1, 3)
>>>
>>> c = a + b
Traceback (most recent call last):
  File "", line 1, in 
ValueError: operands could not be broadcast 
together with shapes (2,3,1,8) (1,3) 
>>>
>>> #  казалось бы что NumPy превратит "b" в 1 x 3 x 1 x 1
... # и все должно получиться вот так:
... #  a: 2 x 3 x 1 x 8
... #  b: 1 x 3 x(1)x(1) 
... #  c: 2 x 3 x 1 x 8
... 
>>> #  Не сработает и такой вариант:
... 
>>> #  a: 2 x 3 x 1 x 8
... #  b:(1)x 1 x 3 x(1) 
... #  c: 2 x 3 x 3 x 8

Именно так и работает первое правило транслирования: если количество осей в массивах не совпадает, то новые оси добавляются только слева. Не справа, не по бокам, а только слева. Причем как вы поняли, иногда достаточно только данного первого правила:

>>> a = np.arange(96).reshape(2, 3, 2, 8)
>>> b = np.arange(16).reshape(2, 8)
>>>
>>> c = a + b
>>> c.shape
(2, 3, 2, 8)
>>>
>>> #  a: 2 x 3 x 2 x 8
... #  b:(1)x(1)x 2 x 8
... #  c: 2 x 3 x 2 x 8
... #     ^---^--------
... #  В массиве "b" добавлены две оси слева

Если количество осей не одинаково, то срабатывает первое правило. А если еще и не одинаковы размеры по некоторым осям, то после первого правила проверяется совместимость размеров и если они оказываются совместимы, то срабатывает второе правило. Мы уже видели как это работает в примере выше, но давайте приведем еще один:

>>> a = np.arange(96).reshape(2, 3, 2, 8)
>>> b = np.arange(2).reshape(2, 1)
>>>
>>> c = a + b
>>> c.shape
(2, 3, 2, 8)
>>>
>>> #  a: 2 x 3 x 2 x 8
... #  b:(1)x(1)x 2 x 1
... #  c: 2 x 3 x 2 x 8

Механизм транслирования чрезвычайно полезная штука, так как арифметические действия над массивами с разными размерами приходится выполнять гораздо чаще, чем можно подумать. Обычно, эти массивы небольшой размерности. Однако, транслирование может сработать даже в случае большого количества осей и избавить вас от написания сложных еще более запутанных, медленных, вложенных циклов.


11.2. Функция ix_()

Мы уже знаем, что благодаря механизму транслирования мы можем выполнять математические операции над массивами разного размера. Однако, для этого их размеры должны быть совместимыми. Допустим у нас есть три массива и мы хотим вычислить новый:

>>> a = np.array([0, 1, 2])
>>> b = np.array([3, 4])
>>> c = np.array([5, 6, 7, 8, 9])
>>>
>>> d = a*b + c    #  Естественно это приведет к ошибке
Traceback (most recent call last):
  File "", line 1, in 
ValueError: operands could not be broadcast together with shapes (3,) (2,)

Логично, что мы все таки можем получить новый массив в котором каждый элемент будет вычисляться по указанным правилам, но для этого размеры массивов a, b и с нужно привести к совместимой форме:

>>> from numpy import newaxis
>>>
>>> a =a[:, newaxis]    #  Добавим две оси справа
>>> a
array([[0],
       [1],
       [2]])
>>>
>>> a = a[:, newaxis]
>>> a
array([[[0]],

       [[1]],

       [[2]]])
>>>
>>> a.shape
(3, 1, 1)
>>>
>>> b = b[:, newaxis]    #  Добавим по одной оси справа и слева
>>> b
array([[3],
       [4]])
>>>
>>> b = b[newaxis, :]
>>> b
array([[[3],
        [4]]])
>>>
>>> b.shape
(1, 2, 1)
>>>
>>> c = c[newaxis, :]    #  Добавим две оси слева
>>> c
array([[5, 6, 7, 8, 9]])
>>>
>>> c = c[newaxis, :]
>>> c
array([[[5, 6, 7, 8, 9]]])
>>>
>>> c.shape
(1, 1, 5)
>>>
>>> d = a*b + c    #  Теперь все массивы совместимой формы
>>> d
array([[[ 5,  6,  7,  8,  9],
        [ 5,  6,  7,  8,  9]],

       [[ 8,  9, 10, 11, 12],
        [ 9, 10, 11, 12, 13]],

       [[11, 12, 13, 14, 15],
        [13, 14, 15, 16, 17]]])

С функцией ix_() все это можно сделать гораздо быстрее:

>>> a = np.array([0, 1, 2])
>>> b = np.array([3, 4])
>>> c = np.array([5, 6, 7, 8, 9])
>>>
>>> an, bn, cn = np.ix_(a, b, c)
>>>
>>> an
array([[[0]],

       [[1]],

       [[2]]])
>>>
>>> bn
array([[[3],
        [4]]])
>>>
>>> cn
array([[[5, 6, 7, 8, 9]]])
>>>
>>> an.shape, bn.shape, cn.shape
((3, 1, 1), (1, 2, 1), (1, 1, 5))
>>>
>>> d = an*bn + cn
>>> d
array([[[ 5,  6,  7,  8,  9],
        [ 5,  6,  7,  8,  9]],

       [[ 8,  9, 10, 11, 12],
        [ 9, 10, 11, 12, 13]],

       [[11, 12, 13, 14, 15],
        [13, 14, 15, 16, 17]]])
>>>
>>> d[2, 0, 4]
15
>>>
>>> a[2]*b[0] + c[4]
15

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

>>> a = np.array([[0, 1],[2, 3]])
>>> b = np.array([4, 5, 6])
>>> c = np.array([[[7, 6],[5, 4]],[[3, 2],[1, 0]]])
>>> a.shape, b.shape, c.shape
((2, 2), (3,), (2, 2, 2))
>>>
>>> an, bn, cn = np.ix_(a, b, c)
Traceback (most recent call last):
  File "", line 1, in 
  File "/home/n1/anaconda3/lib/python3.5/site-packages
   /numpy/lib/index_tricks.py", line 77, in ix_
    raise ValueError("Cross index must be 1 dimensional")
ValueError: Cross index must be 1 dimensional
>>>
>>> an, bn = np.ix_(a, b)
Traceback (most recent call last):
  File "", line 1, in 
  File "/home/n1/anaconda3/lib/python3.5/site-packages
   /numpy/lib/index_tricks.py", line 77, in ix_
    raise ValueError("Cross index must be 1 dimensional")
ValueError: Cross index must be 1 dimensional