📂ジュリア

ジュリアで距離行列計算を最適化する方法

結論

$n$ 個の座標間の距離を計算しようとする。

1. 全ての座標間を計算する必要がなければ、グループに分けて長方形の距離行列を作ればいい。
2. 長方形の距離行列は 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に関しては計算したくもないほど大きな差が出るし、コーディングの難易度もループに比べてずっと簡単だ。

追加研究 ¹

ユークリッド距離が距離関数である場合、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木²を使用しなければならない。NearstNeighbors.jlで速く距離を計算する方法を参照。

環境

• OS: Windows
• julia: v1.5.0