scipy.sparse: 疎行列の要素へ関数の一括適用

くだらないものばかり作ってないでたまには技術ネタを書いてみようと思う。

Python（とくにscipy.sparse）でx=0でf(x)=0となるようなユニバーサル関数fを疎行列のそれぞれの要素に作用させるにはどうすればいいだろうか。

x=0で0になるという点が重要で、その条件を満たす関数ならばもとの行列のsparsityをそのまま保存できる。つまり結果も疎行列になるはずだ。

Pythonにはユニバーサル関数という呼ばれる関数群があって、それは密行列については便利な機能で、例えば、密行列のすべての要素にtanhを作用させようとすると、次のようなコードで計算できる。

import numpy as np

a=np.ones((3,3))
print a
b=np.tanh(a)
print b

つまりユニバーサル関数は行列に作用させると各要素に作用する。以前に書いたfor文を書いたら負けという原則によると、計算スピード的にもかなり有利である。

ところが、aが疎行列のときはnp.tanh(a)とするとエラーになる。これは、一般のユニバーサル関数（引数=0でかならずしも0ではない）を疎行列に作用させると密行列になるので、そういうことをもしやりたいのなら、todense()メソッドで疎行列を密行列に変換してからやればいいということだと思う。

一方でtanhのように引数=0で0になる関数ならば、理屈の上では非ゼロ要素にだけその関数を作用させてあとは放置しておけば正しい結果が得られる。

以下、csr_matrixとcsc_matrixについてだけ考える。実際行列の計算で主に使われるのはcsr_matrixとcsc_matrixだから。解決策についてはcsr_matrixもcsc_matrixも全く同じなので、csr_matrixについてのみ説明する。

前提知識

疎行列の内部形式について説明するが、詳しくは本家のドキュメントを参照。（と、ここで気づいたのだが、内部形式まで触れた日本語の解説は検索しても見つからないので、別途ブログに書いたら需要あるかな？）

csr_matrixのCSRはCompressed Sparse Rowの略で、各行を圧縮した形式。内部データは3つの配列、indices, indptr, dataから構成され、indicesとindptrはインデクスを意味し、dataが要素の値を意味する。i行目（以下、「i行目」や「j列目」といったときのiやjは0からスタートするものとする）の非ゼロ要素の列番号はindices[indptr[i]:indptr[i+1]]に格納され、その各要素はdata[indptr[i]:indptr[i+1]]に格納される。つまり、(i,j)成分がxのとき、indices[indptr[i]:indptr[i+1]]のいずれかがjになり、そのインデクスをkとする（つまりindices[k]==jとなる）と、data[k]にxが格納されている。

csc_matrixのCSCはCompressed Sparse Columnsで、CSRの内部形式の行と列がちょうど逆になった形式である。つまり、j列目の非ゼロ要素の行番号はindices[indptr[j]:indptr[j+1]]に格納され、その要素はdata[indptr[j]:indptr[j+1]]に格納される。

本題

では、x=0でf(x)=0となるユニバーサル関数fを疎行列の各要素に作用させるにはどうすればいいか。つまり数学的に書くと\( A=(a_{ij}) \)が与えられた時に\( B=(f(a_{ij})) \)を計算するにはどうすればいいか。内部データで要素を表すdataの各要素にfを作用させて、indicesとindptrはそのままにすればよい。dataは配列なのでユニバーサル関数の機能をそのまま使えて、csr_matrixのコンストラクタにはindices, indptr, dataを引数にとるやつがあるので、それを使えば効率よく行列が作れる。つまり、fをtanhとすると、

import scipy.sparse as sp
b=sp.csr_matrix((tanh(a.data),a.indices,a.indptr), shape=a.shape)

とすればよい。

ベンチマークを取ってみる。サイズが1万行1万列で、そのうち10万要素だけが非ゼロというケースで試してみる。（Python2で動かしているが、Python3でも、printのところだけ変えれば動くと思う）

import numpy as np
import scipy.sparse as sp
from time import time

N=10000
np.random.seed(0)
s=set()
while len(s)<100000:
    s.add(tuple(np.random.randint(0,N,2)))
a=sp.lil_matrix((N,N))
for i,j in s:
    a[i,j]=1.0
a=a.tocsc()
t0=time()
b0=sp.csc_matrix((np.tanh(a.data),a.indices,a.indptr),shape=a.shape)
t1=time()
print "a) %.3f sec" % (t1-t0)
t0=time()
b1=np.tanh(a.todense())
t1=time()
print "b) %.3f sec" % (t1-t0)
t0=time()
rs,cs=a.nonzero()
b2=sp.lil_matrix(a.shape)
for r,c in zip(rs,cs):
    b2[r,c]=np.tanh(a[r,c])
t1=time()
print "c) %.3f sec" % (t1-t0)

手元のPCでやってみた結果は以下の通り

a) 0.003 sec
b) 0.809 sec
c) 22.305 sec

上から順に、a)前述の手法、b)密行列に変換してからユニバーサル関数の機能を利用、c)非ゼロ要素についてfor文で逐次計算（結果はlil_matrix）、の場合の計算時間だ。結果は一目瞭然だろう。a)が一番速いのは当然だとしても、密行列に変換するよりもfor文の方がはるかに遅いのだから、for文がいかににダメかということがわかる。別途調べたところ22秒のうちlil_matrixへの書き込みが12秒くらいを占めてることがわかったが、それにしてもaへの逐次アクセスは遅い。

まとめ

• 疎行列の非ゼロ要素にまとめてユニバーサル関数を作用させるときは、内部データに手を入れてゴリゴリやるのがいい
• for文でひゼロ要素に逐次実行するくらいなら、密行列に変換してからゴニョゴニョしたが速い

まさかり待ってます♡