Edited at

[Julia] Multiple dispatch の活用 - 行列ベクトル積を例に

More than 1 year has passed since last update.

 Julia では、同じ関数の識別子(名前)に、複数の型の組合せを割り当てられます。呼び出し時の引数の型の組合せで、特定の関数を選んで実行する仕組みが multiple dispatch です。同じ算法(アルゴリズム)を、異なる型のデータに対して実装する際、multiple dispatch を使うと、簡潔に記述(算譜)できます。

 行列・ベクトル積を含む数値計算で、その様子を観察してみます。


問題

 行列 $H$ とベクトル $v_1$ があります。ベクトル列 ${v}_{i}$ を、次の漸化式を用いて発生させるという問題を取り上げます。

\begin{align}

{v}_{2} & = H {v}_{1}, \\
{v}_{i+2} & = H {v}_{i+1} + {v}_{i} \quad i=1,2,3,\ldots
\end{align}

 3つのベクトル $p, q, r$ を用意します。これらを巡回して、漸化式を3回計算します。こうすると、ベクトルをコピーする処理が不要となります。

r = H q + p, \\

q = H p + r, \\
p = H r + q

 初期ベクトル p ($= v_1$) から、ベクトル q($= v_2$) , r ($= v_3$) を確保します。更に、上の2ステップを適用して、最後に得られたベクトルを返す関数 proceed を書きました。参考のため、最後の漸化式計算の時間を計測します。

function proceed(h,p)

q = h * p
r = h * q + p
q = h * p + r
p = h * r + q
#
r = h * q + p
q = h * p + r
@time p = h * r + q
return p
end

 関数 proceed は、引数の型を指定しない "generic" な関数です。

julia> methods(proceed)

# 1 method for generic function "proceed":
proceed(h, p) at REPL[1]:2

 関数 proceed に含まれる計算のうち、積 h * p と和 p + q は、引数の型が未定です。proceed の呼び出し時に hp (の型)が具体的に与えられると、計算が実行される手筈です。


密行列で

 行列 $H$ として、次のような非対称行列を考えます。何も書かれていない要素は零です。非零要素は、対角要素と、上下に隣り合う要素です。

H = \left[

\begin{array}{ccccc}
4 & -2 & & & \\
-1 & 4 & -2 & & \\
& -1 & 4 & \ddots & \\
& & \ddots & \ddots & -2 \\
& & & -1 & 4
\end{array}
\right]

 初期ベクトル p ($= v_1$) は、全要素が 1 のベクトルとします。

 Julia 標準の多次元配列 Array{Float64,N} (N=1, 2) を用いて、hp を確保・計算します。要素数は nn=1000 としました。

nn = 1000

pp0=ones(nn)

hh1=eye(nn)*4.0
for i in 1:nn-1
hh1[i,i+1]=-2.0
hh1[i+1,i]=-1.0
end

 行列ベクトル積 (*)、ベクトル同士の和 (+) は Julia 標準の演算を用いますから、これ以上のプログラムは不要です。では、proceed を呼び出して、実行します。

julia> pp1=copy(pp0)

rr1=proceed(hh1,pp1)
0.000633 seconds (4 allocations: 15.906 KB)
1000-element Array{Float64,1}:
509.0
-214.0
65.0
18.0
21.0

21.0
-27.0
373.0
-919.0
997.0


疎行列で

 零要素を多く含む行列を疎行列 (sparse matrix)といいます。

 Julia には Compressed Sparse Columns (CSC) フォーマットで格納する疎行列 SparseMatrixCSC が実装されています。CSC形式は、あらゆる形の疎行列を扱うことができる優れものですが、一から作るのは面倒かもしれません。

 ここでは、雛形となる密行列 hh1 から、sparse関数を用いて疎行列を生成します。

julia> hh2=sparse(hh1)

1000×1000 sparse matrix with 2998 Float64 nonzero entries:
[1 , 1] = 4.0

 (2998 = nn*3 - 2) 個の非零要素が検出されました。正しく動いていますね。

 SparseMatrixCSC型行列とベクトルの積 (*) は定義済みなので、追加のプログラムは不要です。では、実行。

julia> pp2=copy(pp0)

rr2=proceed(hh2,pp2)
0.000010 seconds (4 allocations: 15.906 KB)
1000-element Array{Float64,1}:
509.0
-214.0
65.0
18.0
21.0

21.0
-27.0
373.0
-919.0
997.0

 最後の漸化式の計算時間が一桁短くなりました。結果が同じであることを確かめましょう。

julia> norm(rr2 - rr1)

0.0


自作の三重対角行列で

 今回の $H$ のように「対角要素と、上下に隣り合う要素が、非零である行列」を、三重対角行列 (tridiagonal matrix)といいます。Julia にも、三重対角行列を扱う Tridiagonal 型が組み込まれていますが、ここでは、三重対角行列型を自作してみましょう。

 MySparseMatrix型は、非零要素を格納する配列 el を内包します。 el の中身は、将来複素数になる可能性があるので、Number 型にしておきます。

abstract MyAbstractSparseMatrix{Tv}

type MySparseMatrix{Tv<:Number} <: MyAbstractSparseMatrix{Tv}
el :: Array{Tv,2}
end

 elの寸法は [n,3] とします。行列 a の対角要素・非対角要素を、以下のように格納することにします。Julia では、第1軸(最初の添字)がメモリ格納順です。

\begin{align}

\ell_{j,1} & = a_{j,j-1} \quad j =2,\ldots, n \\
\ell_{j,2} & = a_{j,j} \quad j =1,\ldots, n \\
\ell_{j,3} & = a_{j,j+1} \quad j =1,\ldots, n-1
\end{align}

 MySparseMatrix型行列 aArray{T,1}型ベクトル b との積を計算する関数 (*) を定義します。 Base.(*) をインポートしてから関数を定義します。

import Base: (*)

function (*){T}(a::MySparseMatrix{T},b::Array{T,1})
n=size(b,1)
c=zeros(b)
c[1]= a.el[1,2]*b[1] + a.el[1,3]*b[1+1]
for i in 2:n-1
c[i]=a.el[i,1]*b[i-1] + a.el[i,2]*b[i] + a.el[i,3]*b[i+1]
end
c[n]=a.el[n,1]*b[n-1] + a.el[n,2]*b[n]
return c
end

 密行列から MySparseMatrix 型の疎行列を生成する関数mySparseMatrixFromDense を作ります。

function mySparseMatrixFromDense(aa)

n=size(aa,1)
el=Array(typeof(aa[1,1]),n,3)
for i in 1:n
el[i,2]=aa[i,i]
end
el[1,3]=aa[1,1+1]
for i in 2:n-1
el[i,1]=aa[i,i-1]
el[i,3]=aa[i,i+1]
end
el[n,1]=aa[n,n-1]
return MySparseMatrix(el)
end

 では、実行してみましょう。

julia> hh3=mySparseMatrixFromDense(hh1)

MySparseMatrix{Float64}([0.0 4.0 -2.0; -1.0 4.0 -2.0; ; -1.0 4.0 -2.0; -1.0 4.0 7.63918e-313])

julia> hh3.el
1000×3 Array{Float64,2}:
0.0 4.0 -2.0
-1.0 4.0 -2.0
-1.0 4.0 -2.0

-1.0 4.0 -2.0
-1.0 4.0 -2.0
-1.0 4.0 7.63918e-313

julia> pp3=copy(pp0)
rr3=proceed(hh3,pp3)
0.000007 seconds (3 allocations: 15.891 KB)
1000-element Array{Float64,1}:
509.0
-214.0
65.0
18.0
21.0

21.0
-27.0
373.0
-919.0
997.0

 最後の漸化式の計算時間が、更に一桁短くなりました。結果が同じであることを確かめましょう。

julia> norm(rr3 - rr1)

0.0


自作の三重対角疎行列を OffsetArraysで

 上の行列・ベクトル積の計算では、添字が 0n+1 となる場合を forループから別扱いしています。もし、これらの場合が結果に影響しないことが保証できるなら、別扱いしないで forループに含めることができます。これは、ベクトル型スーパーコンピュータの算譜で、よく使われた手法です。

 Julia 0.5 以降では、配列添字を 1以外に設定できる仕組みが実装されました。拙文 Qiita記事 「[Julia 0.5] 配列添字を0から数える」 で紹介した OffsetArrays パッケージを用いてみます

 非零要素を格納する配列 elOffsetArray 型にしてみましょう。

using OffsetArrays

immutable MyOffsetSparseMatrix{Tv<:Number} <: MyAbstractSparseMatrix{Tv}
el :: OffsetArray{Tv,2}
end

 elの添字範囲は [0:n+1,-1:1] とします。対角要素・非対角要素を、以下のように格納します。

\begin{align}

\ell_{j,-1} & = a_{j,j-1}, \quad j=2,\ldots, n \\
\ell_{j, 0} & = a_{j,j}, \quad j=1,\ldots, n \\
\ell_{j, 1} & = a_{j,j+1}, \quad j=1,\ldots, n-1
\end{align}

 MyOffsetSparseMatrix型行列 aOffsetArray{T,1}型ベクトル b との積を計算する関数 (*) を定義します。b の添字範囲も [0:n+1] を要求しましょう。行列・ベクトル積は for 文一つで記述できます。行列とベクトルの添字の下限・上限の検査も入れました。

function (*){T}(a::MyOffsetSparseMatrix{T},b::OffsetArray{T,1})

ib=indices(b,1)
@assert ib.start == 0
n=ib.stop-1
ia=indices(a.el,1)
@assert ia.start == 0
@assert ia.stop == n+1
c=zeros(b)
for i in 1:n
c[i]=a.el[i,-1]*b[i-1] + a.el[i,0]*b[i] + a.el[i,1]*b[i+1]
end
return c
end

 密行列から MyOffsetSparseMatrix型疎行列を生成する関数myOffstSparseMatrixFromDense を作ります。

function myOffsetSparseMatrixFromDense(aa)

n=size(aa,1)
el0=zeros(typeof(aa[1,1]),n+2,3)
el=OffsetArray(el0,0:n+1,-1:1)
for i in 1:n
el[i,0]=aa[i,i]
end
el[1, 1]=aa[1,1+1]
for i in 2:n-1
el[i,-1]=aa[i,i-1]
el[i, 1]=aa[i,i+1]
end
el[n,-1]=aa[n,n-1]
return MyOffsetSparseMatrix(el)
end

 

 件の手法がうまくいくためには、el[1,-1]el[n,1] が零である必要があります。myOffsetSparseMatrixFromDense で作成した行列は、これを満たしています。

 ベクトル同士の和は、OffsetArrayで定義済みです。では、実行しましょう。

julia> hh4=myOffsetSparseMatrixFromDense(hh1)

MyOffsetSparseMatrix{Float64}([0.0 0.0 0.0; 0.0 4.0 -2.0; ; -1.0 4.0 0.0; 0.0 0.0 0.0 ; ])

julia> hh4.el
julia> hh4.el
OffsetArrays.OffsetArray{Float64,2,Array{Float64,2}} with indices 0:1001×-1:1:
0.0 0.0 0.0
0.0 4.0 -2.0
-1.0 4.0 -2.0

-1.0 4.0 -2.0
-1.0 4.0 0.0
0.0 0.0 0.0

julia> pp40=zeros(nn+2)
pp4=OffsetArray(pp40,0:nn+1)
pp4[1:nn]=pp0

julia> pp4
OffsetArrays.OffsetArray{Float64,1,Array{Float64,1}} with indices 0:1001:
0.0
1.0

1.0
0.0

julia> rr4=proceed(hh4,pp4)
0.001290 seconds (25.95 k allocations: 655.906 KB)
OffsetArrays.OffsetArray{Float64,1,Array{Float64,1}} with indices 0:1001:
0.0
509.0
-214.0
65.0
18.0
21.0

21.0
-27.0
373.0
-919.0
997.0
0.0

 記述は簡潔になりましたが、計算時間は最悪となりました。メモリ確保が多数回起こっているのも、気になります。

 結果が同じであることを確かめましょう。ベクトルの添字は [0:nn+1] ですが、計算結果が入っている添字範囲は [1:nn] であることに注意します。

julia> norm([rr4[i] - rr1[i] for i in 1:nn])

0.0


展望

 行列ベクトル積を例に multiple dispatch の様子を観察しました。

 以上では割愛しましたが、今回の例に含まれる $r = h * q + p$ の計算は、BLASxGEMV 関数 ($y \leftarrow \alpha A x + \beta y$) を用いても実装できそうです。 Juliaは、BLASのラッパーを含んでおり gemv!関数を直接呼び出せます。機会を見て紹介したいです。

 行列の積という基本演算でも、計算機性能を引き出すには、そのアーキテクチャや命令体系を考慮したチューニングが必須です。(参考 → Qiita記事 「プロセッサの性能を引き出すのは大変だよというお話 (行列積チューニング)」Common Lisp で行列計算するブログを追試してみた」「 同 (その2)」「同 (その3)」および記事内リンク、特に「片桐, プログラム高速化の基礎」) 。実際の数値計算は、基本演算を豊富に含みますから、これらを取り替えて検討するのに multiple dispatch の仕組みが役立つでしょう。 私自身は Templates for the Solution of Linear Systems: Building Blocks for Iterative Method とか、Templates for the Solution of Algebraic Eigenvalue Problems: a Practical Guide などの文書を眺めて、色々試したくて、うずうずしています。