Optimization question

I’m trying to optimize the code block below:

using Distributions
order = 2
σ_Priors = Array{Distribution,2}(undef,order,order)
ϵ_Priors = Array{Distribution,2}(undef,order,order)
σ_Priors[:,:] .= Gamma(10,0.1)
ϵ_Priors[:,:] .= Gamma(10,0.1)
std_Prior = Gamma(10,0.3)

function totalEnergy(crystal::Crystal,LJ::LJ)
    totalEnergy = 0
    for i in eachindex(LJ.ϵ)
        totalEnergy += -LJ.ϵ[i] * LJ.σ[i] ^6 * crystal.r6[i] + LJ.ϵ[i] * LJ.σ[i] .^12 * crystal.r12[i]
    end
#        @time "total Energy" totalEnergy = sum(-LJ.ϵ .* LJ.σ .^6 .* crystal.r6 .+ LJ.ϵ .* LJ.σ .^12 .* crystal.r12 ) 
end

function logNormal(data::MatSim.DataSet,LJ::MatSim.LJ,σ::Float64)
    thesum = 0
    for i = 1:data.nData
        thesum += (data.crystals[i].energyFP -MatSim.totalEnergy(data.crystals[i],LJ))^2
    end
    thesum *= -data.nData/2 *log(σ^2) - 1/(2 * σ^2)
#    @time "thesum" thesum =  -data.nData/2 *log(σ^2) - 1/(2 * σ^2) * sum([(i.energyFP - totalEnergy(i,LJ))^2   for i in data.crystals])
    return thesum

end

function logPost(data,model,σ)
    # return 3.0
    thesum = logNormal(data,model,σ)  + logpdf(std_Prior,σ)
    for i in eachindex(model.σ)
        thesum += logpdf(σ_Priors[i],model.σ[i])
        thesum += logpdf(ϵ_Priors[i],model.ϵ[i])
    end
    return thesum
end

@time logPost(trainingSet,LJ,0.3)

I recognize that I’m not giving you all of the code but I’m hoping it’s enough to go on. Currently a single call to logPost requires 34 allocations. Doesn’t seem too bad except that this function is being called hundreds of thousands of times. I’m convinced that this function is the performance bottleneck. Can anyone see any glaring problems here? Perhaps using globals for σ_Priors instead of passing them into the function?

Thanks in advance.

Yea, mark that const is a good start

Or if I just pass them into the function that would work too right?

1 Like

Distribution is an abstract type, hence

logpdf(σ_Priors[i],model.σ[i])

is a dynamic (and thus expensive) dispatch. If you only have gamma distributions, do this instead

σ_Priors = fill(Gamma(10,0.1), order, order)

(same for ϵ_Priors)

3 Likes

For example, following the previous tips of the other thread:

With a parametric struct:

julia> struct ModelA{T}
           σ::Vector{Float64}
           ϵ::Vector{Float64}
           σ_Priors::Matrix{T}
           ϵ_Priors::Matrix{T}
       end

julia> function logPost(model)
           # return 3.0
           thesum = 0.0 # simplified version
           for i in eachindex(model.σ)
               thesum += logpdf(model.σ_Priors[i],model.σ[i])
               thesum += logpdf(model.ϵ_Priors[i],model.ϵ[i])
           end
           return thesum
       end
logPost (generic function with 2 methods)

julia> m = ModelA(rand(100), rand(100), fill(Gamma(10,0.1), 10, 10), fill(Gamma(10,0.1), 10, 10));

julia> typeof(m)
ModelA{Gamma{Float64}}

julia> @btime logPost($m)
  3.100 μs (0 allocations: 0 bytes)
-675.8052525561492

With the abstract field, instead:

julia> struct ModelB
           σ::Vector{Float64}
           ϵ::Vector{Float64}
           σ_Priors::Matrix{Distribution}
           ϵ_Priors::Matrix{Distribution}
       end

julia> mB = ModelB(m.σ, m.ϵ, m.σ_Priors, m.ϵ_Priors);

julia> typeof(mB)
ModelB

julia> @btime logPost($mB)
  6.455 μs (600 allocations: 9.38 KiB)
-675.8052525561492

FWIW, I think this blog post might be helpful: New blog post about Julia parametric types and constructors

I think perhaps this is not causing any problems right here, but I don’t think it’s good style to shadow the typename like this.

1 Like

I ran this exact code and got


0.000010 seconds (1 allocation: 16 bytes)

Note that it’s @btime not @time (from BenchmarkTools.jl).

1 Like

I see what you’re saying. This declaration:

σ_Priors = Array{Distribution,2}(undef,order,order)

is problematic. It turns out that not all of the distributions are identical. They may all be Gammas, but with different parameters. How could I avoid this problem in this case?

If you know all possible Distributions then you can just pack them into a single type by using DynamicalSumTypes.jl

Does the gamma distribution use different types for different parameters? If they are all the same type, you could just use typeof(Gamma(10, 0.1)) for example. If you have a few different distribution types, wrapping them in a Union type could be fine. Alternatively, use LightSumTypes.jl as mentioned above

1 Like

I’m showing that the function logPost is type unstable. Since model.σ_Priors[i] is a distribution the type of logpdf is Any.

Here’s the snippet from @code_warntype:

%138 = thesum::Any
%139 = Base.getproperty(metrop, :ϵ_Priors)::Matrix{Distribution}
%140 = i@_29::Int64
%141 = Base.getindex(%139, %140, j)::Distribution
%142 = Base.getproperty(model, :ϵ)::LinearAlgebra.UpperTriangular{Float64, Matrix{Float64}}
%143 = i@_29::Int64
%144 = Base.getindex(%142, %143, j)::Float64
%145 = MatSim.logpdf(%141, %144)::Any
%146 = (%138 + %145)::Any

My model.σ is an upper triangular matrix:


σ:: UpperTriangular{Float64, Matrix{Float64}}
ϵ:: UpperTriangular{Float64, Matrix{Float64}}

Could that be part of the problem?

I don’t think so. Where that UpperTriangular type is coming from?

The main reason is probably the abstract Distribution type. If you only need distributions of type Gamma as elements of of the priors (or other of the same type with varying parameters), you should parameterize the type.