📂줄리아

줄리아에서 거리 행렬 계산 최적화하는 법

결론

$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