Juliaの行列ベクトル積について(疎行列編)
科学技術計算では行列ベクトル積が頻繁に登場し,これを効率よく計算することはとても重要です.
行列サイズ $n$ に対して計算量は $\mathrm{O}(n^2)$ ですが,行列は疎行列(ほとんどの要素が $0$)であることが多く,定義通りに計算すると $0$ を掛けたり足したりする無駄な演算が多く発生してしまいます.
以下では疎行列とベクトルの積の計算について解説します.
(もっとも,Juliaで行列ベクトル積 $A\boldsymbol{b}$ を計算するには A*b
と書けばよく,これは $A$ が疎行列の場合でも同様なので,疎行列のパッケージがあることさえ知っていれば実用上は十分なことが多いのですが,とはいえ,チューニングを行ったり,上三角部分など行列の一部とベクトルの積を計算したりしたいときには,疎行列とベクトルの積がどのように計算されているか理解しておくことが重要です.)
まず,行列 $A$ を適当に定義して,疎行列に変換してみましょう.
julia> using SparseArrays
julia> A = [1 2 5; 0 3 0; 0 4 6]
3×3 Array{Int64,2}:
1 2 5
0 3 0
0 4 6
julia> A = sparse(A)
3×3 SparseMatrixCSC{Int64,Int64} with 6 stored entries:
[1, 1] = 1
[1, 2] = 2
[2, 2] = 3
[3, 2] = 4
[1, 3] = 5
[3, 3] = 6
マニュアルによれば,圧縮列格納方式 (CSC) が採用されています.
struct SparseMatrixCSC{Tv,Ti<:Integer} <: AbstractSparseMatrix{Tv,Ti}
m::Int # Number of rows
n::Int # Number of columns
colptr::Vector{Ti} # Column i is in colptr[i]:(colptr[i+1]-1)
rowval::Vector{Ti} # Row indices of stored values
nzval::Vector{Tv} # Stored values, typically nonzeros
end
上で定義した行列に対しては,それぞれ以下のようになっています.
julia> A.m
3
julia> A.n
3
julia> A.colptr
4-element Array{Int64,1}:
1
2
5
7
julia> A.rowval
6-element Array{Int64,1}:
1
1
2
3
1
3
julia> A.nzval
6-element Array{Int64,1}:
1
2
3
4
5
6
さて,行列ベクトル積 $y=Ax$ を計算するためには,定義の通り,
$$y_i = \sum_{j=1}^nA_{ij}x_j$$
をfor文で実装するのが自然ですが,
- 疎行列の場合,$0$にアクセスするのは無駄
- 行列の要素は行ではなく列を基準に格納されている
ため,少し工夫する必要があります.例えば,次のように書くのがよいでしょう.
x = ones(A.n);
y = zeros(A.n);
for col = 1 : A.n
idx = A.colptr[col]
for i = idx : (A.colptr[col + 1] - 1)
y[A.rowval[i]] += A.nzval[i] * x[col]
end
end
3重対角行列でテストしてみます.
using SparseArrays
using LinearAlgebra #Tridiagonalを使うために必要
function main()
function matvec(A,x)
y = zeros(A.n);
for col = 1 : A.n
idx = A.colptr[col]
for i = idx : (A.colptr[col + 1] - 1)
y[A.rowval[i]] += A.nzval[i] * x[col]
end
end
return y
end
n = 400000; #行列サイズ
A = Tridiagonal(ones(n-1),-2ones(n),ones(n-1)); #三重対角行列
A = sparse(A); #疎行列に
# あるいは最初から
# A = spdiagm(1=>ones(n-1),0=>-2*ones(n),-1=>ones(n-1))
x = ones(n);
@time matvec(A,x); #時間測定
end
main()
計算時間は行列サイズ $n$ にほぼ比例するはずです(行列サイズが2倍になれば計算時間も約2倍).