scipy.sparseでの疎ベクトルの扱い

Pythonで疎行列を扱うライブラリであるscipy.sparseについて、自分がハマってきたところをまとめてみようと思う。とくに、行列よりもベクトルとして扱ったとき（つまり行または列の数が1のとき）の注意点をまとめる。

基本的なことばかりなのかもしれないが、日本語の情報は少ないので、それでも役に立てるかもと思いました。まあ自分の勉強ノートみたいなもんです。

基本

scipy.sparseで提供される疎行列の形式はいくつかあるが、ここではとくにCOO(COOdinate)形式、LIL(LInked List)形式、CSR(Compressed Sparse Rows)形式、CSC(Compressed Sparse Columns)形式について取り上げる。

それらの違いについて、詳細は本家のドキュメントを参照して欲しいが、すごく大雑把にまとめると、

• CSR, CSC形式は、同じ型同士の掛け算・足し算が高速。LIL、COO形式は演算には向いていない。
• CSRは行を取り出すのが高速、CSCは列を取り出すのが高速
• LIL形式は[i,j]成分に値をセットするのは高速。CSR、CSC形式は、[i,j]成分に値をセットするのは遅い（warningが出る）。COOは[i,j]成分に値をセットする機能そのものがない。
• 以上のようなことがあるので、使い勝手を考えると、LIL形式疎行列を作って値をセットする→CSRまたはCSCに変換して計算をする、という流れが一般的（ただし、以下に書くようにこれには罠がある）

ノルムの計算

まず本題に入る前に、以下のソースコードこのようなヘッダを共通で暗黙に使うことにする。

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

疎行列\(v\)（列行列とする）のノルム\(|v|^2\)を計算したいとする。そのとき、素直に考えて\(v^T v\)を計算しようとすると効率がよくない。\(v\)がCSC形式だとすると、\(v^T\)はCSR形式になるからだ。なのでそういうときは要素ごと掛け算multiplyを使うほうがいい。

実際にやってみる。

n=100000000
ii=np.arange(0,n,2)
ij=(ii,np.zeros(len(ii),dtype=int))
data=np.array([1.0]*len(ii))
a=sp.csc_matrix((data,ij),shape=(n,1))
t0=time()
x=a.T.dot(a)[0,0]
print "%.3f sec" % (time()-t0)
t0=time()
y=a.multiply(a).sum()
print "%.3f sec" % (time()-t0)
assert abs(x-y)<1e-7

そして実験結果。

0.778 sec
0.329 sec

multiplyを使ったほうが倍くらい速い。

これは一般の内積にも言えることで、型違いの掛け算するよりは要素積を使ったほうがいい。

疎行列の初期化

疎行列というのはほとんどの要素がゼロなので、非ゼロ要素だけの情報を持てばよいことになっている。なので、全要素ゼロの行列を取得すること自体はほとんどコストがかからないような気がするかもしれない。しかし、それが罠だったりする。

LIL, COO, CSC, CSR形式について、それぞれ列ベクトル、行ベクトルを作って、消費されるメモリ、時間を測ってみる。

def matsize(a):
    if isinstance(a,sp.lil_matrix):
        return a.data.nbytes+a.rows.nbytes
    elif isinstance(a,sp.coo_matrix):
        return a.data.nbytes+a.row.nbytes+a.col.nbytes        
    else:
        return a.data.nbytes+a.indices.nbytes+a.indptr.nbytes

n=100000000
# case 1
t0=time()
a=sp.lil_matrix((n,1))
print "1: %d,  %.3f sec" % (matsize(a),time()-t0)
a=None
# case 2
t0=time()
a=sp.lil_matrix((1,n))
print "2: %d,  %.3f sec" % (matsize(a),time()-t0)
a=None
# case 3
t0=time()
a=sp.coo_matrix((n,1))
print "3: %d,  %.3f sec" % (matsize(a),time()-t0)
a=None
# case 4http://docs.scipy.org/doc/scipy/reference/sparse.html
t0=time()
a=sp.coo_matrix((1,n))
print "4: %d,  %.3f sec" % (matsize(a),time()-t0)
a=None
# case 5
t0=time()
a=sp.csc_matrix((n,1))
print "5: %d,  %.3f sec" % (matsize(a),time()-t0)
a=None
# case 6
t0=time()
a=sp.csc_matrix((1,n))
print "6: %d,  %.3f sec" % (matsize(a),time()-t0)
a=None
# case 7
t0=time()
a=sp.csr_matrix((n,1))
print "7: %d,  %.3f sec" % (matsize(a),time()-t0)
a=None
# case 8
t0=time()
a=sp.csr_matrix((1,n))
print "8: %d,  %.3f sec" % (matsize(a),time()-t0)

（ちなみに、ここでmatsize関数は、オブジェクトのサイズを正確に測っているのではなくて、行列を表現するための主な要素のサイズをバイト数で計算している）

そして実験結果がこちら。

1: 1600000000,  90.860 sec
2: 16,  0.000 sec
3: 0,  0.000 sec
4: 0,  0.000 sec
5: 8,  0.000 sec
6: 400000004,  0.112 sec
7: 400000004,  0.113 sec
8: 8,  0.000 sec

時間に関してはLIL＆列ベクトルが圧倒的にかかっている。メモリについては、LIL＆列ベクトル、CSC＆行ベクトル、CSR＆列ベクトルのコストが高い。データの内部表現（詳細は本家ドキュメント参照）を考えればなぜそうなるかは自明かとは思うが、特にLIL＆列ベクトルのときについてのみ説明すると、これは、空リストn個からなる配列を2つ用意している。つまり2*n個の空リストを用意しているので重い。ここでは計測してないが、小さいオブジェクトをたくさん用意してるので、開放時のガーベジコレクションも重いので注意が必要だ（サンプルコード中、次の処理の前にNoneを代入しているのはそういう理由）。

これが要注意なのは、よくあるscipy.sparseのサンプルコードではlil_matrixを使っているからだ。確かにlil_matrixで用意してa[i,j]=vのようなコードを書くのは便利なのだけど、全体のサイズが大きいとメモリと時間を余計に食うことになる。大きい行列を最初に1個だけ用意してそれを何度も使うような計算なら初期化コストは気にならないかもしれないが、大きい行列をたくさん用意する場合だと致命的になる。

ではどうすればいいかというと、a[i,j]=vのような使い方は諦めて最初からcsc_matrixやcsr_matrixを使うようにして、__init__の引数で全部の情報を渡してやればいい。

ちなみにlil_matrixのドキュメントには大きい行列には向いてないって書いてあるのだけど、ハマるまで気づかなかった。そもそも疎行列って大きい行列を扱うためのものじゃないのかと小一時間くらい問い詰めたい。

要素が0と1だけの疎ベクトルとの内積

ショッピングのレコメンドのアルゴリズムとか考えていると、よく要素が1と0だけの（つまり非ゼロ要素はすべて1の）疎行列を考えることがあるのだが、最初から0/1行列だってわかっていれば非ゼロはすべて1なのだから、掛け算を省いて高速化できないかと考えてみる。

n=100000000
ii1=np.arange(0,n,100)
ij1=(ii1,np.zeros(len(ii1),dtype=int))
ii2=np.arange(0,n,2)
ij2=(ii2,np.zeros(len(ii2),dtype=int))
data1=np.array([1.0]*len(ii1))
data2=np.array([0.5]*len(ii2))
a=sp.csc_matrix((data1,ij1),shape=(n,1))
b=sp.csc_matrix((data2,ij2),shape=(n,1))
t0=time()
x=a.multiply(b).sum()
print "%.3f sec" % (time()-t0)
t0=time()
ii,_=a.nonzero()
y=b[ii,0].sum()
print "%.3f sec" % (time()-t0)
assert abs(x-y)<1e-7

ここでは、インデクス配列を利用したインデクシングを使っている。

そして実験結果がこちら。

0.147 sec
0.977 sec

余計な小細工しないで素直に掛け算したほうが速かった様子。あらら。

まとめ

疎行列のライブラリscipy.sparseは、便利な機能があってありがたいのだけれど、本当にブラックボックスとして使うと痛い目にあうかもしれない（私は実際に痛い目にあった）。詳細まで知る必要はないにしても（それがライブラリ使うことの良さなんだから）、内部データ構造やアルゴリズムの概要くらいは理解してないと使いこなせいと思う。つまり、ドキュメント読み込むこと大事。