I have some code which loads a csv file of 2000 2D coordinates, then a function called collision_count
counts the number of pairs of coordinates that are closer than a distance d
of each other:
using BenchmarkTools
using CSV
using LinearAlgebra
function load_csv()::Array{Float64,2}
df = CSV.read("pos.csv", header=0)
return Matrix(df)'
end
function collision_count(pos::Array{Float64,2}, d::Float64)::Int64
count::Int64 = 0
N::Int64 = size(pos, 2)
for i in 1:N
for j in (i+1):N
@views dist = norm(pos[:,i] - pos[:,j])
count += dist < d
end
end
return count
end
结果如下:
pos = load_csv()
@benchmark collision_count($pos, 2.0)
BenchmarkTools.Trial:
memory estimate: 366.03 MiB
allocs estimate: 5997000
--------------
minimum time: 152.070 ms (18.80% GC)
median time: 158.915 ms (20.60% GC)
mean time: 158.751 ms (20.61% GC)
maximum time: 181.726 ms (21.98% GC)
--------------
samples: 32
evals/sample: 1
这比此Python代码慢30倍:
import numpy as np
import scipy.spatial.distance
pos = np.loadtxt('pos.csv',delimiter=',')
def collision_count(pos, d):
pdist = scipy.spatial.distance.pdist(pos)
return np.count_nonzero(pdist < d)
%timeit collision_count(pos, 2)
5.41 ms ± 63 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
有什么方法可以使其更快?那所有的分配又如何呢?
我能得到的最快的速度如下
这里有各种各样的变化,有些是风格上的,有些是结构上的。从样式开始,您可能会注意到,我输入的内容没有限制。这没有任何性能优势,因为Julia足够聪明,可以为您的代码推断类型。
The biggest structural change is switching from using a matrix to a vector of
StaticVectors
. The reason for this change is that since points are your scalar type, it makes more sense to have a vector of elements where each element is a point. The next change I made is to use a squared norm, sincesqrt
operations are expensive. The results speak for themselves:Note that there are
n log(n)
algorithms that may be faster, but this should be pretty close to optimal for a naive implementation.