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])
            # 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]
            # 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])
            for l in 1:3
                coords[l, i] = coord_new_i[l]
                coords[l, j] = coord_new_j[l]
        # Decrease learning rate after each cycle
        learning_rate -= 1 / no_cycles
    return coords

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]))

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.