python numpyで3×3行列とベクトルの配列にたいして内積を行いたい

Question

前提・実現したいこと

python3.7 numpyを使用して、行列の配列に対して内積をしたいと考えております。
行列は、
A[3][3][x]
B[3][x]
xは数千程度で、
A[3][3]・B[3]の結果をx分格納した結果を取得したいのですが、
numpyのdot積分を使用して、
np.dot(A,B.T)
と記述すると、思った結果が得られないという状況です。

例として、次のようなコードを書いた場合に
a = np.asarray([[[1, 2 , 3], [4, 5, 6], [7, 8, 9]]
, [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
, [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
, [[1, 2, 3], [4, 5, 6], [7, 8, 9]]])
b = np.asarray([[1, 0, 0],[2, 0, 0], [3, 0, 0], [4, 0, 0]])
c = np.dot(a, b.T)
print(c)
求める結果は
[[ 1 4 7]
[ 2 8 14]
[ 3 12 21]
[ 4 16 28]]
なのですが、実際の出力は、
[[[ 1 2 3 4]
[ 4 8 12 16]
[ 7 14 21 28]]

[[ 1 2 3 4]
[ 4 8 12 16]
[ 7 14 21 28]]

[[ 1 2 3 4]
[ 4 8 12 16]
[ 7 14 21 28]]

[[ 1 2 3 4]
[ 4 8 12 16]
[ 7 14 21 28]]]
となります。

これを解決するために、
c = np.empty([b.shape[0],b.shape[1]])
for i in range(b.shape[0]):
buf = b[i]
c[i] = np.dot(a[i],buf.T)
print(c)
としたところ、
[[ 1. 4. 7.]
[ 2. 8. 14.]
[ 3. 12. 21.]
[ 4. 16. 28.]]
となり、求める結果を得ることはできました。

しかし、せっかくnumpyを使用しているのに、for文で回しているので、
for文を使わずに、もっといい書き方があれば、ご教授いただければと思います。

宜しくお願いいたします。

Answer 1

a = np.array([[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
              [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
              [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
              [[1, 2, 3], [4, 5, 6], [7, 8, 9]]])
b = np.array([[1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0]])

c = (a*b[:,None,:]).sum(2)
print(c)
# [[ 1  4  7]
#  [ 2  8 14]
#  [ 3 12 21]
#  [ 4 16 28]]

c = np.einsum('ijk,ik->ij', a, b)
print(c)
# [[ 1  4  7]
#  [ 2  8 14]
#  [ 3 12 21]
#  [ 4 16 28]]

ということですか？

英語ですが、以下のサイトが参考になります。
A basic introduction to NumPy's einsum – ajcr – Haphazard investigations

---

"""
i = 第0軸
j = 第1軸
k = 第2軸
"""
----- a -----
    j 0       1       2
i   k 0 1 2
0 #[[[1 2 3] [4 5 6] [7 8 9]]
1 # [[1 2 3] [4 5 6] [7 8 9]]
2 # [[1 2 3] [4 5 6] [7 8 9]]
3 # [[1 2 3] [4 5 6] [7 8 9]]]

----- b -----
    j 0       1       2
i   k 0 1 2
0 #[ [1 0 0]
1 #  [2 0 0]
2 #  [3 0 0]
3 #  [4 0 0]                 ]

"""
A{i, j, k} と B{i, j, k} を要素ごとに掛け算したい。
Bには [i,j,k] の3軸のうち、真ん中の軸が欠けている。
→→Bの2軸の間にj軸を追加すると以下のようになる
"""
----- b[:,None,:] -----
    j 0       1       2
i   k 0 1 2
0 #[[[1 0 0] [  >  ] [  >  ]]
1 # [[2 0 0] [  >  ] [  >  ]]
2 # [[3 0 0] [  >  ] [  >  ]]
3 # [[4 0 0] [  >  ] [  >  ]]]

"""
k（第2軸）を総計して潰すので、sum(axis=2)
"""
----- (a*b[:, None, :]).sum(2) -----
    j  0    1    2
i   k  -
0 # [[ 1    4    7]
1 #  [ 2    8   14]
2 #  [ 3   12   21]
3 #  [ 4   16   28]]

Comments

[kyte_tsuboyama]

j軸がかけている分を追加する必要があったのですね。
とても参考になりました、ありがとうございました。
アインシュタインの縮約記法はちょっと難しかったので、
勉強してみたいと思います。

Answer 2

 x = map(lambda x: list(np.dot(a,x)[0]), b)
 print(list(x))

速度が向上するかわかりませんが、例えばmapを利用する例です。