优化O(n²)复杂度的Julia代码，提升仿真计算效率

背景 #

我正在开发一款时间复杂度为O(n²)的仿真程序，主体基于Python，为提升性能将核心计算卸载至Julia。每次仿真迭代需对比100个实例，共产生4900次比较：Python实现耗时约63秒，Julia初始实现耗时降至30秒。这是我首次编写Julia代码，认为仍有优化空间，现寻求改进方案。

calculate_novelty由Python脚本调用，histograms是包含100个20×20矩阵的向量。

初始Julia代码 #

using Shuffle

function apply_noise_mask(hist, INT_CASTER::UInt)
    diff::UInt = INT_CASTER - sum(hist)

    mask = zeros(UInt, size(hist))
    mask = shuffle(replace(mask, 0=>1, count=diff))
    return hist + mask
end

function move_supply(supply, fcoord::CartesianIndex , capacity, tcoord::CartesianIndex, INT_CASTER::UInt)
    if supply[fcoord] <= capacity[tcoord]
        flow = supply[fcoord]
        capacity[tcoord] -= flow
        supply[fcoord]= 0
    else
        flow = capacity[tcoord]
        supply[fcoord] -= flow
        capacity[tcoord] = 0
    end
    distance = sqrt(abs(fcoord[1]-tcoord[1])^2 + abs(fcoord[2]-tcoord[2])^2)
    score = flow/INT_CASTER*distance
    return score, supply, capacity
end

function wasserstein_distance(hist0, hist1, INT_CASTER::UInt)::Float16
    supply, capacity = trunc.(UInt, hist0.*INT_CASTER), trunc.(UInt, hist1.*INT_CASTER)
    supply, capacity = apply_noise_mask(supply, INT_CASTER), apply_noise_mask(capacity, INT_CASTER)

    score::Float16 = 0
    while true
        from_idx, to_idx = findfirst(supply .> 0), findfirst(capacity .> 0)

        if (from_idx == nothing) || (to_idx == nothing)
            return score
        end
        work, supply, capacity = move_supply(supply, from_idx, capacity, to_idx, INT_CASTER)
        score += work
    end
end

function calculate_novelty(histograms)
    @time begin

        INT_CASTER::UInt = 10000
        amt_instances::UInt = length(histograms) # population = 100 -> int8 goes until 127

        novelty_scores = zeros(Float16, amt_instances)
        for i = 1:amt_instances-1
            for j = 1+i:amt_instances
                score = wasserstein_distance(histograms[i], histograms[j], INT_CASTER)
                novelty_scores[i] += score
                novelty_scores[j] += score
            end
        end
    end
    return novelty_scores
end

初始运行性能 #

15.760257 seconds (20.87 M allocations: 28.901 GiB, 11.97% gc time)
15.592724 seconds (20.72 M allocations: 28.697 GiB, 11.95% gc time)

Julia版本：1.8

已完成的优化及效果 #

替换findfirst为自定义遍历函数 #

function find_first_candidate(arr)
    for elem in CartesianIndices(arr)
       if arr[elem] > 0
          return elem
        end
    end
    return nothing
end

优化后平均运行耗时：

1.095999 seconds (99.02 k allocations: 314.210 MiB, 1.74% gc time)

采用原地修改函数 #

进一步减少内存分配后，平均运行耗时：

1.090733 seconds (59.41 k allocations: 188.528 MiB, 1.70% gc time)

寻求进一步优化方案 #

现希望得到更多针对该Julia代码的优化建议，以进一步提升性能。

内容的提问来源于stack exchange，提问作者Oliver