ゼロが大半の係数行列は、行番号・列番号・値の組からcoo_arrayで組み立て、計算用途に合わせてCSRなどへ変換する。重複した座標の値は合算できるので、局所的な寄与を積み上げる処理に向いている。ただし疎配列と従来の疎行列では演算子の意味が異なるため、型を確認して使おう。
最小例で確かめる
以下は説明用に作った小さなデータである。実測データや実行速度の測定結果ではない。コード全体をexample.pyとして保存すれば、入力ファイルを別途用意せずに実行できる。assertは、この例で成り立つべき形や値を確認するために入れてある。
import numpy as np
from scipy.sparse import coo_array, csr_array
# 4点の隣接結合。対角(0,0)には二つの寄与を足す。
row = np.array([0,0,1,2,3,0,1,1,2,2,3])
col = np.array([0,0,1,2,3,1,0,2,1,3,2])
data = np.array([1.,1.,2.,2.,2.,-1.,-1.,-1.,-1.,-1.,-1.])
coo = coo_array((data, (row, col)), shape=(4,4))
print("stored contributions:", coo.nnz)
coo.sum_duplicates()
a = coo.tocsr()
print("after combining duplicates:", a.nnz)
print("small dense check:")
print(a.toarray()) # この4×4の例だけを密配列にする。
x = np.array([1.,2.,3.,4.])
print("matrix-vector product:", (a @ x).tolist())
print("elementwise square diagonal:", (a * a).diagonal().tolist())
assert isinstance(a, csr_array)
expected = np.diag(np.full(4,2.)) + np.diag(np.full(3,-1.),1) + np.diag(np.full(3,-1.),-1)
np.testing.assert_array_equal(a.toarray(), expected)
np.testing.assert_allclose(a @ x, [0.,0.,0.,5.])
np.testing.assert_array_equal((a * a).toarray(), expected * expected)
np.testing.assert_array_equal((a @ a).toarray(), expected @ expected)
assert a.nnz == 10
n = 100_000
print("dense float64 payload [GB]:", n*n*8 / 10**9)
print("CSR payload [bytes]:", a.data.nbytes + a.indices.nbytes + a.indptr.nbytes)
実行結果
stored contributions: 11
after combining duplicates: 10
small dense check:
[[ 2. -1. 0. 0.]
[-1. 2. -1. 0.]
[ 0. -1. 2. -1.]
[ 0. 0. -1. 2.]]
matrix-vector product: [0.0, 0.0, 0.0, 5.0]
elementwise square diagonal: [4.0, 4.0, 4.0, 4.0]
dense float64 payload [GB]: 80.0
CSR payload [bytes]: 200
COOは座標と値の一覧から作る
row、col、dataの同じ位置の要素が、一つの行列成分への寄与を表す。例では対角成分が2、隣同士の結合が-1の4×4配列を作っている。(0, 0)だけは1を二つ登録し、局所的な寄与を後で足し合わせる状況を再現した。shapeを指定することで、末尾のゼロ行・列も含めた全体の大きさが分かる。
COO形式は重複座標を保持できるので、作成直後のnnzは11である。sum_duplicatesで同じ座標を合算すると10になる。nnzは保存されている要素数であり、数学的に非ゼロである成分数と常に同じとは限らない。明示的な0が保存されている場合には、その0も含まれる。
組み立てと計算で形式を選ぶ
行方向の操作や行列ベクトル積などにはCSRが扱いやすいので、例ではtocsrで変換した。COO、CSR、CSCなどは保存方法が違い、得意な操作も異なる。どの形式でもすべてのindex操作が同じように使えるわけではないため、組み立てと計算の段階で適した形式を選ぶ。
データの値を変えずに形式を変える操作と、同じ座標の寄与を足す操作は意味が違う。形式変換の際に重複が合算される場合もあるので、どの段階で何件になるか把握しておく。ゼロが明示的に残っている場合にはeliminate_zerosも使えるが、構造として0の位置を保持したい用途では削除の意味を考える必要がある。
疎配列の*と@を混同しない
この例で使うcoo_arrayとcsr_arrayは疎配列である。*は要素ごとの積、@は行列積を表す。従来のcoo_matrixやcsr_matrixでは*の意味が異なるため、古いコードを型だけ置き換えると計算内容が変わるおそれがある。行列積には@を使うと意図を明示しやすい。
例ではa * aとa @ aを別々の密配列による計算と比較した。またa @ xは[0, 0, 0, 5]になることを検査している。小さな既知の配列で演算の意味を確認してから大きくすれば、shapeが合っているだけの誤計算を見つけやすい。ndarray向けの関数へ疎配列を渡す際も対応状況を確認しよう。
密配列化する前にサイズを見積もる
10万×10万のfloat64密配列は、値の領域だけで80GBになる。コードで示したのは要素数×8byteの算術的な見積もりであり、その大きな配列を実際に確保したり速度測定したりしてはいない。toarrayを呼ぶと疎に保存した利点を失うので、デバッグ表示のつもりでもサイズを確認する。
CSRでは値以外に列添字や行の区切りを保存する領域も必要である。例のnbytesの和はそれらの配列領域の合計で、Pythonオブジェクト全体の使用メモリ測定ではない。小さい行列やゼロが少ない行列では疎形式の負担が相対的に大きくなる。疎なら必ず速い・省メモリという結論ではなく、非ゼロの割合と行う演算を合わせて判断しよう。
動作確認環境と参考資料
Linux・CPython 3.12.14、NumPy 2.3.5、pandas 2.2.3、SciPy 1.17.0、Matplotlib 3.10.8の環境で掲載コードを実行した。使用するライブラリはコード冒頭のimportを参照してほしい。公式資料の最新版と、この実行確認版は区別している。数値の末尾や表の表示幅は環境によって変わることがある。