Just curiosity, I implemented a (non-general, not very optimized) form of the cell lists above. It resulted to be be ~8 times slower than the method implemented in NearestNeighbors.jl
. I will have to check if the method there applied to my actual use case . I hope it does, having it 8 times faster won’t be bad at all.
NN(A, B) = (0.00044736197216615207, 1080399, 1136)
0.918342 seconds (7.54 k allocations: 96.740 MiB)
MD(A, B) = (0.000447361972166152, 1080399, 1136)
7.259169 seconds (4 allocations: 11.452 MiB, 0.11% gc time)
(0.000447361972166152, 1080399, 1136)
import Pkg
Pkg.activate("nn")
using BenchmarkTools
using NearestNeighbors, Random, Printf
N1 = 15 * 10^5
N2 = 15 * 10^2
# create input data:
Random.seed!(1); A = rand(3, N1)
Random.seed!(1234); B = rand(3, N2)
# NearestNeighbors algorithm:
function NN(A,B)
kdtree = KDTree(A)
idxs, dists = nn(kdtree, B) # convenience function for k=1
ixb = argmin(dists)
ixa = idxs[ixb]
return dists[ixb], ixa, ixb
end
@show NN(A,B)
@time NN(A,B)
function icell3D(x, n)
i = trunc(Int,(x[1]/n))+1
j = trunc(Int,(x[2]/n))+1
k = trunc(Int,(x[3]/n))+1
return i, j, k
end
function MD(A,B)
d = 0.1 # assume that we know that the shortest distance is smaller than this
n = floor(Int,1/d)
# cell lists
fatomA = zeros(Int,n,n,n) # first atom of A in the cell
natomA = zeros(Int,size(A,2)) # next atom of A in the cell
# fill cells
for i in axes(A,2)
ic, jc, kc = icell3D(@view(A[:,i]),n)
natomA[i] = fatomA[ic,jc,kc]
fatomA[ic,jc,kc] = i
end
iamin = 0
ibmin = 0
dmin = +Inf
# loop over atoms of B
@inbounds for ib in axes(B,2)
xb = @view(B[:,ib])
i, j, k = icell3D(xb,n)
# Loop over vicinal cells to compute distances to solvent atoms, and
# add data to dc structure (includes current cell)
for ic in i-1:i+1
(ic == 0 || ic > n) && continue
for jc in j-1:j+1
(jc == 0 || jc > n) && continue
for kc in k-1:k+1
(kc == 0 || kc > n) && continue
ia = fatomA[ic,jc,kc]
while ia > 0
dij = (A[1,ia]-B[1,ib])^2 +
(A[2,ia]-B[2,ib])^2 +
(A[3,ia]-B[3,ib])^2
if dij < dmin
dmin = dij
iamin = ia
ibmin = ib
end
ia = natomA[ia]
end
end
end
end
end
return sqrt(dmin), iamin, ibmin
end
@show MD(A,B)
@time MD(A,B)