Imagine you have a series of n points randomly generated in a box in 3D space. You also have a list of distance bounds, e.g. points 5 and 3 should be between 1.0 and 2.0 Angstroms apart.
There are various methods to get from random coordinates to a solution that fills some or all distance bounds. Below is an implementation of Stochastic Proximity Embedding, one such method, in Julia. I am looking for input on improving the speed.
function run_spe(present_i, present_j, bounds_lower, bounds_upper, no_atoms)
squares_lower = bounds_lower .* bounds_lower
squares_upper = bounds_upper .* bounds_upper
learning_rate = 1.0
coord_new_i = zeros(3)
coord_new_j = zeros(3)
bound_near = 0.0
no_cycles = 1500
no_steps = 70000
no_bounds = length(present_i)
# Generate random starting coordinates
coords = rand!(zeros(3, no_atoms)) * 100.0
for cycle in 1:no_cycles
for step in 1:no_steps
# Pick a random bound from the list of bounds
rand_ind = rand(1:no_bounds)
i = present_i[rand_ind]
j = present_j[rand_ind]
# Compute square distance between atoms
sq_dist = 0.0
for m = 1:3
sq_dist += abs2(coords[m, i] - coords[m, j])
end
# Check if square distance is within bounds
# If it is, go to the next iteration
# If it is not, find the nearest bound
if sq_dist < squares_lower[rand_ind]
bound_near = bounds_lower[rand_ind]
elseif sq_dist > squares_upper[rand_ind]
bound_near = bounds_upper[rand_ind]
else
continue
end
# Adjust coordinates to bring them closer to the nearest bound
distance = sqrt(sq_dist)
prefactor = learning_rate * 0.5 * (bound_near - distance) / distance
for k in 1:3
coord_new_i[k] = coords[k, i] + prefactor * (coords[k, i] - coords[k, j])
coord_new_j[k] = coords[k, j] + prefactor * (coords[k, j] - coords[k, i])
end
for l in 1:3
coords[l, i] = coord_new_i[l]
coords[l, j] = coord_new_j[l]
end
end
# Decrease learning rate after each cycle
learning_rate -= 1 / no_cycles
end
return coords
end
present_i
/present_j
are lists of indices. bounds_lower
/bounds_upper
are lists of distances. For example, present_i[10] = 5
, present_j[10] = 3
, bounds_lower[10] = 1.0
, bounds_upper[10] = 2.0
would mean the tenth bound corresponds to the example bound above.
I have done a bit of profiling and the major bottlenecks are the initial lines in the inner loop (from rand_ind up to continue). Ideally Julia tricks to speed individual steps up would be helpful, but if there is a better data structure than the lists I use that would also be good.
Sample data can be found here for testing and read in using something quick like:
# Set in_filepath as path to downloaded and uncompressed file
present_i = Int[]
present_j = Int[]
bounds_lower = Float64[]
bounds_upper = Float64[]
no_atoms = 2035
open(in_filepath, "r") do in_file
for line in eachline(in_file)
cols = split(strip(line), '\t')
push!(present_i, int(cols[1]))
push!(present_j, int(cols[2]))
push!(bounds_lower, float(cols[3]))
push!(bounds_upper, float(cols[4]))
end
end