Juliaの行列ベクトル積について（疎行列編）

Feb 2, 2022 2 min read

科学技術計算では行列ベクトル積が頻繁に登場し，これを効率よく計算することはとても重要です．行列サイズ $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倍）．

Julia

Yuto Miyatake

Associate Professor

Applied Mathematics, Numerical Analysis, Computational Uncertain Quantification