Ruby で高速に多次元数値配列を扱いたいときに Numo::NArray を使っているのだが、この配列は基本的に「要素毎」の演算となっている。では NArray で行列の積などを計算したいときはどうすればいいのか気になったので、調べて整理した。
TL;DR
- よく使う積はメソッドが用意されていて、行列の積なら
a.dot(b)でいい - 自分でやる場合、行列の積 $C_{ij} = \sum_{k} A_{ik} B_{kj}$ は
c = a[false, :new].mulsum(b, -2)で計算できる（a,bは2次元配列） - より一般にテンソルの縮約を計算する場合、掛け合わせるテンソルの添字が揃うように配列の次元を追加し、
#mulsumで縮約する次元を指定して積和計算する
（前提知識）ビューの次元指定について
a[false, :new, true] と書かれていても意味がわからないときのためのメモ。
従って以下は全て同じビューを得られる。（一部の方法は なお、特別な意味を持つ引数は他にもあるみたいだが（該当コード）、ドキュメントなどでの使用例が少ないので、本記事では クリックして詳細を表示
a = Numo::Int64.new(2, 4, 6).indgen に対して a[1, 2..3, 0..-1] のように次元毎にインデックスや範囲を指定して部分配列を切り出せる1
a.ndim と合っていないといけない（例なら3つ）
…というのはあくまで基本ルールで、 false や :new は別途考える必要がある0..-1 など全範囲を表すものは true と略記できるtrue は、1ヶ所だけ false で置き換えられる（rest dimensions）
:new を追加すると、そこに大きさ1の次元を追加できる
#expand_dims というメソッドでも可能（ただし1回につき1ヶ所のみ）a.ndim == 3 を前提としている）# a = Numo::Int64.new(2, 4, 6).indgen
a[false, :new, true] #=> Numo::Int64(view)#shape=[2,4,1,6] [...]
a[false, true].expand_dims(-2)
a[true, true, true].expand_dims(-2)
a[0..-1, 0..-1, 0..-1].expand_dims(-2)
a.view.expand_dims(-2)
a.expand_dims(-2)
true, false, :new で統一する。
機能
引数に与えるオブジェクト
全範囲
true, nil, :all, :*
全範囲・逆順
:reverse
全範囲・集計フラグ設定
:reduce, :sum, :+
残余次元
false, :rest, :~
次元追加
:new, :-
1次元配列同士の積
手始めに、ベクトルに見立てた2つの1次元配列の演算を扱う。
require 'numo/narray'
u = Numo::Int64[1, 2, 3, 4]
v = Numo::Int64[1, 10, 100, 1000]
単純に配列同士を掛け算すると、要素毎の積（アダマール積）が返る。
u * v
#=>
# Numo::Int64#shape=[4]
# [1, 20, 300, 4000]
内積（ドット積）を計算するには、要素毎の積の総和をとればいい： $a = \sum_i u_i v_i$ 。これは前述の計算の後に #sum を作用させてもいいし、 #mulsum で一気に計算することもできる。
(u * v).sum #=> 4321
u.mulsum(v) #=> 4321
もうひとつ別の積として外積（二項積、直積、テンソル積）2も考えられる： $A_{ij} = u_i v_j$ 。これは九九の表を作るイメージでよく、ベクトルの大きさは異なっていてもいい。 NArray で実現するには、大きさ1の次元を追加して添字が被らないようにして積をとる3。
# i, j
u[false, :new].shape #=> [4, 1]
v .shape #=> [4]
u[false, :new] * v
#=>
# Numo::Int64#shape=[4,4]
# [[1, 10, 100, 1000],
# [2, 20, 200, 2000],
# [3, 30, 300, 3000],
# [4, 40, 400, 4000]]
行列の積の計算では、この内積と外積の考え方を組み合わせる。
2次元配列による行列の積
行列の積は $C_{ij} = \sum_{k} A_{ik} B_{kj}$ と計算する。図で説明しているもの4を見るとわかりやすいが、 C の各要素は A の行と B の列との内積となっている（また、これを全ての行と列の組に対して計算しているのは外積と似ている）。内積は前節で扱ったように積の総和という形に分解して考えていい。
NArray は基本的に要素毎の積を計算するので、行と列というように次元がずれているものは計算できない。そこで配列に大きさ1の次元を追加して、内積をとりたい次元を揃えることで解決する。行列の積であれば、上式 A の末尾に次元をひとつ追加すると k は後ろから2次元目で揃う（添字が早く回るのは後ろ側なので、次元は後ろから数えて合わせること）。あとは内積計算を完了するために k の次元で総和をとればいい。
ちなみに k の次元を揃えたことで i, j の次元は相手側の大きさが1になっている。これは前節の外積と同じで、全ての (i, j) の組について計算することを表す。
require 'numo/narray'
a = Numo::Int64[[1, 2], [3, 4]]
b = Numo::Int64[[1, 10], [100, 1000]]
# i, k, j
a[false, :new].shape #=> [2, 2, 1]
b .shape #=> [2, 2]
a[false, :new] * b # kの次元を揃えた状態で要素毎の積を計算する
#=>
# Numo::Int64#shape=[2,2,2]
# [[[1, 10],
# [200, 2000]],
# [[3, 30],
# [400, 4000]]]
(a[false, :new] * b).sum(-2) # 後ろから2番目(k)の次元で和をとる
#=>
# Numo::Int64#shape=[2,2]
# [[201, 2010],
# [403, 4030]]
a[false, :new].mulsum(b, -2)
#=> (same)
積と和を分けて計算すると、積をとった段階で3次元配列が作られメモリを多く消費する。そのため #mulsum で一気に計算したほうが効率がいい。
一般化
数式が添字表記で、どの次元で和をとるのかわかっていれば、同様に次元追加（および必要なら転置）して #mulsum で計算すればいい。物理学で式をテンソルで表す場合などによく登場し、総和記号を省略することもある。
例えば $b_{pq} = \sum_{i,j} r_{pi} r_{qj} a_{ij}$ （2階テンソルの座標変換をイメージした式）なら、要素を表す添字は $pq, pi, qj, ij$ が登場している。順序を保つように全ての添字を並べる方法は [p, q, i, j] または [p, i, q, j] があり、これに合うよう次元追加すれば転置不要で計算できる。（実装時は、最初から全ての添字を揃える必要は無く、二項演算の度に次元追加しても構わない）
r = Numo::DFloat[[0.8, 0.6], [-0.6, 0.8]]
a = Numo::DFloat[[1, 2], [3, 4]]
# p, q, i, j
r[true, :new, true, :new].shape #=> [3, 1, 3, 1]
r[ true, :new, true].shape #=> [3, 1, 3]
a .shape #=> [3, 3]
# 数式通りの順序だと、途中で4次元配列ができてしまう
b = (r[true, :new, true, :new] * r[true, :new, true]).mulsum(a, -2, -1)
# （要素の演算が可換・結合的なら）縮約しやすい順序に変えて構わない
# ij q i j j p q i i
b = a.mulsum(r[true, :new, true], -1).mulsum(r[true, :new, true], -1)
# p, i, q, j
r[true, true, :new, :new].shape #=> [3, 3, 1, 1]
r .shape #=> [3, 3]
a[ true, :new, true].shape #=> [3, 1, 3]
b = a[true, :new, true].mulsum(r[true, true, :new, :new], -3).mulsum(r, -1)
一方で、この式は行列の積でも書き表せるが、掛ける順序に制約ができるうえに転置が必要になる： $b_{pq} = \sum_{i,j} r_{pi} a_{ij} r^T_{jq}$ 。
class Numo::NArray
def matmul(other)
self[false, :new].mulsum(other, -2)
end
end
b = r.matmul(a).matmul(r.transpose(-1, -2))
3階以上でも問題ない。例えば $\sigma_{ij} = \sum_{k,l} d_{ijkl} \epsilon_{kl}$ （弾性係数テンソルをイメージした式）なら [i, j, k, l] で合わせればいい。1回の乗算で2つ以上の添字を縮約するときは、 #mulsum に指定する次元を増やす5。
d = Numo::Int64.new(3, 3, 3, 3).rand(10)
e = Numo::Int64.new(3, 3).rand(10)
# i, j, k, l
d.shape #=> [3, 3, 3, 3]
e.shape #=> [3, 3]
s = d.mulsum(e, -2, -1)
レヴィ＝チヴィタ記号で遊んでみたりも。 $\sum_{j,k} \epsilon_{ijk} \epsilon_{ljk} = 2 \delta_{il}$ が成り立つか試してみる。
e = Numo::Int8.zeros(3, 3, 3)
e[0, 1, 2] = e[1, 2, 0] = e[2, 0, 1] = 1
e[0, 2, 1] = e[1, 0, 2] = e[2, 1, 0] = -1
# indices: i, l, j, k
e.expand_dims(-3).mulsum(e, -2, -1)
#=>
# Numo::Int8#shape=[3,3]
# [[2, 0, 0],
# [0, 2, 0],
# [0, 0, 2]]
定義済みの積
よく使う積に関しては、自分で実装しなくても Numo::NArray のメソッドとして存在している。内積を計算する箇所は実際に #mulsum を利用している。
- ドキュメント： https://ruby-numo.github.io/narray/narray/Numo/NArray.html
- ソースコード： https://github.com/ruby-numo/numo-narray/blob/v0.9.0.7/lib/numo/narray/extra.rb
-
#dot： ドット積／行列積 - 多次元配列の後ろ2次元を行列とみなし（1次元配列なら行or列ベクトル）、要素（行列）毎の積を計算する。なので普通の行列計算ならこのメソッドに任せるのが良さそう。
なお、線形代数ライブラリ Numo::Linalg が利用可能な場合はそちらへ投げている。 -
#inner： 内積 - 多次元配列をベクトルの配列とみなし、ベクトル毎の内積を計算する。何次元目がベクトルかはキーワード引数で指定可能（デフォルトは -1 、つまり最後の次元）。
-
#outer： 外積（直積） - 多次元配列をベクトルの配列とみなし、ベクトル毎の外積を計算する。内積は次元が減る（ベクトル→スカラー）のに対し、こちらは次元が増える（ベクトル→2階テンソル）。
-
#kron： クロネッカー積 - テンソル積を計算し、各要素を多次元配列内に展開する。
それぞれ試してみた。
require 'numo/narray'
a = Numo::Int64[[1, 2], [3, 4]]
b = Numo::Int64[[1, 10], [100, 1000]]
a.dot(b)
#=>
# Numo::Int64#shape=[2,2]
# [[201, 2010],
# [403, 4030]]
a.inner(b)
#=>
# Numo::Int64#shape=[2]
# [21, 4300]
# 参考：行列の積を内積で再現
a.expand_dims(-1).inner(b.expand_dims(-3), axis: -2)
#=>
# Numo::Int64#shape=[2,2]
# [[201, 2010],
# [403, 4030]]
a.outer(b)
#=>
# Numo::Int64#shape=[2,2,2]
# [[[1, 10],
# [2, 20]],
# [[300, 3000],
# [400, 4000]]]
a.kron(b)
#=>
# Numo::Int64#shape=[4,4]
# [[1, 10, 2, 20],
# [100, 1000, 200, 2000],
# [3, 30, 4, 40],
# [300, 3000, 400, 4000]]
# 参考：テンソル積
a[false, *Array.new(b.ndim, :new)] * b
#=>
# Numo::Int64#shape=[2,2,2,2]
# [[[[1, 10],
# [100, 1000]],
# [[2, 20],
# [200, 2000]]],
# [[[3, 30],
# [300, 3000]],
# [[4, 40],
# [400, 4000]]]]