ジュリアで距離行列計算を最適化する方法
結論
$n$ 個の座標間の距離を計算しようとする。
- 全ての座標間を計算する必要がなければ、グループに分けて長方形の距離行列を作ればいい。
- 長方形の距離行列は
pairwise()
関数で簡単かつ速く計算できる。
速度比較
例えば、SIRモデルに対して移動するエージェントベースのシミュレーションを行うと考えてみよう。元の時間計算量は $O \left( n^{2} \right)$ だが、$S$ と $I$ のグループに分けて計算すると、時間計算量は $O \left( n_{S} n_{I} \right)$ に大きく下がる。通常、病気の伝播は $S$ と $I$ 間の距離行列を計算して一定の半径 $\varepsilon$ 内に入るかを判断し、どれだけ接触したかによって実装される。この作業で速度を比較してみよう。
using Distances
using StatsBase
using Random
Random.seed!(0);
N = 10^4
location = rand(2, N);
state = sample(['S', 'I'], Weights([0.1, 0.9]), N);
S = location[:, state .== 'S']
I = location[:, state .== 'I']
function foo(S, I)
contact = Int64[]
for s in 1:996
push!(contact, sum(sum((I .- S[:,s]).^2, dims = 1) .< 0.1))
end
return contact
end
@time foo(S, I)
もちろん、グループに分けて計算するという発想は、ジュリアだけでなくどんな手法を使っても性能向上をもたらす。ポイントは、無理にループを回す必要がなく、Distance
パッケージの pairwise()
関数を上手く使えばいいということだ。
julia> @time foo(S, I);
0.170835 seconds (7.98 k allocations: 210.854 MiB, 12.56% gc Time)
julia> @time sum(pairwise(Euclidean(),S,I) .< 0.1, dims = 1);
0.087875 seconds (14 allocations: 69.639 MiB, 4.15% gc Time)
これら二つの命令は正確に同じ機能を持つが、速度面では約2倍の差があり、allocationに関しては計算したくもないほど大きな差が出るし、コーディングの難易度もループに比べてずっと簡単だ。
追加研究 1
ユークリッド距離が距離関数である場合、Euclidean()
の代わりに SqEuclidean()
を使うと、ルートを取る計算を省略できるため、速度がさらに上がる。次は正確に同じ結果を出すが、速度面では約1.5倍の差が出る。
julia> @time sum(pairwise(Euclidean(),S,I) .< 0.1, dims = 1);
0.091917 seconds (14 allocations: 69.639 MiB, 7.60% gc Time)
julia> @time sum(pairwise(SqEuclidean(),S,I) .< 0.01, dims = 1);
0.061776 seconds (14 allocations: 69.639 MiB, 11.37% gc Time)
さらに、もっと速くなることができる。ここでは、単純なコード最適化だけでは速度を上げるのが難しく、多次元検索に有利なデータ構造であるk-d木2を使用しなければならない。NearstNeighbors.jlで速く距離を計算する方法を参照。
環境
- OS: Windows
- julia: v1.5.0