NumPyの使い方（１５） ブロードキャスト

NumPyのブロードキャストについて。
　
配列の次元、大きさ（要素数）などが異なっていても配列どうしやスカラーの演算ができる。
ブロードキャストができるのは、2つの配列のshapeを最終軸から順に比較していき、同じ大きさか、一方が1であるかチェック。
１つでも該当しない軸があるとブロードキャストできない。
　

　
例をいくつか。

import numpy as np

a1 = np.array([[1, 2, 3], [4, 5, 6]])
print(a1)     # shape:(2, 3)
# [[1 2 3]
#  [4 5 6]]

print(a1 + 1)
# [[2 3 4]
#  [5 6 7]]

print(a1 + [1, 1, 1]) # shape:(3,)
# [[2 3 4]
#  [5 6 7]]

#print(a1 + [1, 1])
# ValueError: operands could not be broadcast together with shapes (2,3) (2,) 

print(a1 + [[1], [10]]) # shape:(2, 1)
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

a1 = np.arange(0, 600, 100).reshape(2, 3, 1)
print(a1)
# [[[0]
#   [100]
#   [200]]
#
#  [[300]
#   [400]
#   [500]]]

a2 = np.arange(12).reshape(3, 4)
print(a2)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

print(a1 + a2)
# [[[  0   1   2   3]
#   [104 105 106 107]
#   [208 209 210 211]]
#
#  [[300 301 302 303]
#   [404 405 406 407]
#   [508 509 510 511]]]

　
ブロードキャスト可能かどうか

b = (lambda a1, a2:all(i==1 or j==1 or i==j for i, j 
          in zip(a1.shape[::-1], a2.shape[::-1])))

a1 = np.arange(6)

print(b(a1.reshape(2, 3, 1), a1.reshape(3, 2))) # True
print(b(a1.reshape(2, 3, 1), a1.reshape(2, 3))) # False
print(b(a1.reshape(2, 3, 1), a1.reshape(1, 6))) # True
print(b(a1.reshape(2, 3, 1), a1.reshape(6, 1))) # False
print(b(a1.reshape(2, 1, 3), a1.reshape(2, 3))) # True
print(b(a1.reshape(3, 1, 2), a1.reshape(6,)))   # False
print(b(a1.reshape(2, 1, 3), np.zeros(3)))      # True
print(b(a1.reshape(2, 1, 3), np.array(10)))     # True スカラー(0次元)