how can I make this faster (just want to learn Julia better)?

Viewed 78

I am writing a small MD (molecular dynamics) engine in Julia and I noticed that 67% of the time spent running the following Ewald summation function is spent in garbage collection. Besides algorithmic improvements, which can certainly be made here, what can I do to make it faster in terms of Julia tricks / good practices?

function getEwaldEnergy(top::Topology, xyz::Array{Float64})::Float64
    # cubic box only for now
    tolerance::Float64 = 1.0e-24
    tryUntil::Int64 = 4  # 5 box sizes gets it to five decimal places
    L::Float64 = top.box[1]
    rc::Float64 = L / 2.0

    α::Float64 = 3.5 / (top.box[1] / 2)
    σ::Float64 = 1 / (sqrt(2) * α)

    λ::Float64 = 1.0
    ns::Array{Int64} = getCombinations(-tryUntil, tryUntil, 3)

    # self-energy corrections
    selfEnergy::Float64 = 0.0
    for i::Int64 in 1:top.nAtoms
        selfEnergy += (1 / (sqrt(2 * pi) * σ)) * (top.charges[i] * top.charges[i])
    end

    # short-range energy
    shortRangeEnergy::Float64 = 0.0
    shortRangeEnergyBefore::Float64 = 0.0
    for l::Int64 in 1:size(ns)[1]
        shortRangeEnergyBefore = shortRangeEnergy
        for i in 1:top.nAtoms
            for j::Int64 in 1:top.nAtoms
                ci::Float64 = top.charges[i]
                cj::Float64 = top.charges[j]

                idx::Array{Array{Int64}} = top.bondIdx
                bonded::Bool = false

                # check if i/j pair is bonded
                for k::Int64 in 1:length(idx)
                    if (i == idx[k][1] && j == idx[k][2]) || (j == idx[k][1] && i == idx[k][2])
                        bonded = true
                        break
                    end
                end

                if bonded == false && !(i == j && ns[l, :] == [0.0, 0.0, 0.0])
                    r::Float64 = norm(xyz[i, :] - xyz[j, :] + L * ns[l, :])
                    if r < rc - λ
                        shortRangeEnergy += 0.5 * ((ci * cj) / r) * erfc(r / (sqrt(2) * σ))
                    elseif r < rc && r > rc - λ
                        # switching function turns on
                        shortRangeEnergy += 0.5 * ((ci * cj) * (rc - r)^2 * (-2 * rc + 2 * r + 3 * λ) * erfc(r / (sqrt(2) * σ))) / (2 * r * λ^3)
                    end
                end
            end
        end
    end


    # long-range energy
    longRangeEnergy::Float64 = 0.0
    longRangeEnergyBefore::Float64 = 0.0
    volume::Float64 = L^3
    ms = ns

    longRangeEnergyBefore = copy(longRangeEnergy)
    for l in 1:size(ms)[1]  # no k = [0, 0, 0]
        if ms[l, :] == [0.0, 0.0, 0.0]
            continue
        end

        # define reciprocal space vectors
        kVector::Array{Float64} = (2 * pi) * (ms[l, :] / L)
        k::Float64 = norm(kVector)

        # check this equal 1 good
        # println(exp(-im * dot(kVector, L * [ms[l, 1], ms[l, 2], ms[l, 3]])))
        strucFactor::ComplexF64 = 0.0
        for i in 1:top.nAtoms
            strucFactor += top.charges[i] * exp(im * dot(kVector, xyz[i, :]))
        end

        longRangeEnergy += ((4 * pi) / (volume)) * ((exp(-σ^2 * k^2 / 2) / k^2) * (abs(strucFactor)^2))
    end


    # this roughly cancels out long range energy due to bonded interactions (Tuckerman pg. 663)
    bondEnergy::Float64 = 0.0
    idx = top.bondIdx
    for b::Int64 in 1:size(idx)[1]
        for i in 1:top.nAtoms
            for j in 1:top.nAtoms
                if ((idx[b][1] == i && idx[b][2] == j) || (idx[b][1] == j && idx[b][2] == i)) # && !(i in isdone)
                    # println("firing j: ", j)
                    ci::Float64 = top.charges[i]
                    cj::Float64 = top.charges[j]
                    # rx::Float64 = xyz[i, 1] - xyz[j, 1]
                    # ry::Float64 = xyz[i, 2] - xyz[j, 2]
                    # rz::Float64 = xyz[i, 3] - xyz[j, 3]
                    rij::Float64 = norm(xyz[i, :] - xyz[j, :])
                    bondEnergy += (ci * cj * erf(rij / (sqrt(2) * σ))) / rij
                end
            end
        end
    end

    ewaldEnergy = shortRangeEnergy + longRangeEnergy - selfEnergy - bondEnergy

    return ewaldEnergy
end
0 Answers
Related