【dask】daskで線形代数の計算をしてみた

## はじめに
daskでお手軽に線形代数計算の並列処理ができるということで試してみた。

環境は以下の通り。

+ macOS Mojave 10.14.6
+ python 3.7.3
+ dask 1.1.4

実行はvscodeだったりjupyterlabだったりする。

## daskって？
daskの基本的な使い方などについては以下の記事が大変参考になる。

+ [PythonのDaskをしっかり調べてみた（大きなデータセットを快適に扱う）](https://qiita.com/simonritchie/items/e174f243bc03fb25462e)

ただし、線形代数関係の細かい部分については自分で調べる必要があった。特に、[公式ページのAPI一覧](https://docs.dask.org/en/latest/array-api.html)に載っていない関数とかを使いたい場合は[GitHubのコード](https://github.com/dask/dask/blob/master/dask/array/linalg.py)を見た方が早かったりする。

## やりたいこと
やりたいのは[以前の記事](https://qiita.com/shachah-svaahaa/items/f622a0e22dbba92e36d2)で扱った行列の低ランク近似である。最初は自分で書いたわけだが、daskに実装されていると知り、そっちの方が速かろうといいうことで試してみた次第である。
低ランク近似に関する細かい話は上記の記事をご参照いただければ幸いである。


## 使い方
早速使ってみる。といっても、難しいことは何もなく、以下のサンプルコードを見てもらえばわかると思う。

```python
import numpy as np
import dask.array as da

# テストデータ
n = 10**6
m = 100
data = np.random.randn(n,m)

D = da.from_array(data, chunks='auto')
Qt = da.linalg.compression_matrix(D, q=30, n_power_iter=0) 
B = Qt @ D

Q = (Qt.T).compute()
B = B.compute()
```

乱数行列に対して低ランク近似は精度が出ないと思うが、そこはただの実行例ということでご寛恕いただきたい。

### 直交行列のサイズと転置
上記のコードのメインは近似直交行列を計算する以下の部分である。

```python
Qt = da.linalg.compression_matrix(D, q=30, n_power_iter=0)
```

主なパラメータはターゲットランクとべき乗反復の回数だが、オーバーサンプリングを明示的に指定することはできない。内部的には`compression_level`という関数でオーバーサンプリングは$10$に設定されている。まぁ、パラメータといいつつ、この値以外に設定する特別な理由もないので、外からいじれないようにされても特に問題はない。

注意点として、近似直交行列は$\mathbf{Q}^{\top}$という転置された形が戻り値になっている。そのため、最初、自分で転置してから行列$\mathbf{B}$を計算しようとしたところ、`Chunks do not add up to shape.`というエラーになった（サイズが合わないので行列積が計算できない）。

### 計算時間とべき乗反復

簡単に`time`モジュールで計算時間を計測してみたところ、以下のようになった（平均はしてない）。

||べき乗反復|計算時間 [s]|
|:--|:-:|:-:|
|`compression_matrix`|0|3.31|
||2|4.63|
||4|5.89|
|||
|自作コード|0|5.17|
||2|15.33|
||4|25.45|

自分で書いたナイーブな実装と比べると、当然のごとくdaskの方が速い。しかも、自作コードはべき乗反復の回数を増やすと計算時間がかなりかかったのに対し、daskの方は計算時間の増加は非常に緩やかだった。実際の運用ではさらに大規模になる上、データの読み込み（HDF5）も絡むので、余計にdaskの方が優秀であろうことが分かる。実にすばら。

## TSQR分解

もう1つ試してみる。daskにはTSQR分解が実装されている。これは縦長行列（tall and skinny matrix）、すなわち、行数よりも列数の方がはるかに小さい行列に対して並列効率のよいQR分解の手法である。これは元の行列を行方向に分割し、個々のブロック行列をQR分解した結果を最後にまとめるという仕方で求められる（アルゴリズムの詳細などは[ここ](https://ipsj.ixsq.nii.ac.jp/ej/?action=repository_action_common_download&item_id=98707&item_no=1&attribute_id=1&file_no=1)とか[ここ](http://www.kurims.kyoto-u.ac.jp/~kyodo/kokyuroku/contents/pdf/1791-07.pdf)に書いてある）。縦長行列であれば通常のQR分解よりも高速化できるため、上述した`compression_matrix`も内部のQR分解にはこの方法を用いている。

```python
# QR分解
Q, R = qr(data, mode='economic')

# TSQR分解
d = da.from_array(data, chunks='auto')
Q, R = da.linalg.tsqr(d, compute_svd=False)
Q = Q.compute()
R = R.compute()
```

上記と同じ乱数行列に対してjupyterの`%%timeit`で計算時間を見てみると、以下の通り。

||計算時間 [s]|
|:--|:-:|:-:|
|`da.linalg.tsqr`|3.81|
|`da.linalg.tsqr`（$\mathbf{Q}$のみ）|2.17|
|`scipy.linalg.qr`|4.76|

特に、直交行列$\mathbf{Q}$のみを計算した場合は`scipy.linalg.qr`の半分以下の時間で済んだ。
QR分解の目的がベクトルや行列の直交化である場合は$\mathbf{R}$は不要なので、そういったケースではかなりの高速化が見込めると思われる。

とはいえ、通常のQR分解より計算量が増えるらしいので、正方行列だったりメモリに乗っかるサイズだったりする場合にはそこまで高速ではないだろう。場面に応じて通常のQR分解と使い分けることで、トータルの計算時間をより削減できると思う。

## まとめ

daskで行列計算、かなりよさげな感じだった。もちろん、現状、numpy/scipyの線形代数関係のライブラリすべてに対応しているわけではないので、何でもかんでもdaskでということでもないだろうが、適用できそうな状況なら試してみる価値は大いにあると思う。

現状、daskより速いと噂のvaexには残念ながら線形代数関係のライブラリが入ってない様子なので、自分にとっては今後もしばらくはdaskの方が有用なツールになりそう。