Question about program design with ComponentArrays

I’m retaking a program I did some time ago, trying to speed it up. I am using ComponentArrays because it is very easy to organize data and parameters with it. Consider this:

using ComponentArrays
using Parameters: @unpack

const homecode = 6

function agecat(a)
    (19 <= a <= 24) && return 1
    (25 <= a <= 29) && return 2
    (30 <= a <= 34) && return 3
    (35 <= a <= 39) && return 4
    (40 <= a <= 44) && return 5
    (45 <= a <= 49) && return 6
    (50 <= a <= 54) && return 7
    (55 <= a <= 59) && return 8
    (60 <= a <= 64) && return 9
end

const data = ComponentArray(
    wẘ = rand(2, 9, 2, 2, 5, 44),
    P_1 = rand(5, 44),
    ē = rand(5, 44)
)

const θ = ComponentArray(
    α0 = rand(2, 2, 5),
    α = rand(3),
    α̃h = rand(2, 2),
    α̃ga = rand(2, 9)
)

function Uk(g, a, e, k, t; data=data, θ=θ)
    @unpack α0, α = θ
    @unpack wẘ, P_1, ē = data
    w = wẘ[:, :, :, 1, :, :]

    α0[g,e,k] + α[1] * w[g, agecat(a), e, k, t] + α[2] * (e - ē[k, t])^2 + α[3] * P_1[k, t]
end

function α̃0(g, a, e, t; data=data, θ=θ)
    @unpack α̃h, α̃ga = θ

    α̃h[1,g] + α̃h[2,g] * t + α̃ga[g, agecat(a)]
end

function U0(g, a, e, t; data=data, θ=θ)

    α̃0(g, a, e, t; data=data, θ=θ)
end

function Probk(g,a,e,k,t; data=data, θ=θ)

    uk = ( k == homecode ? U0(g, a, e, t; data=data, θ=θ) : Uk(g, a, e, k, t; data=data, θ=θ) )
    1/(sum(exp(Uk(g, a, e, kk, t; data=data, θ=θ)-uk) for kk=1:5) + exp(U0(g, a, e, t; data=data, θ=θ)-uk))
end

Notice how I pass data and θ from Probk to U0 and Uk. Is that good practice? I’m worried about allocations.

Ultimately, I will have to optimize a function over θ and I am trying squeeze as much speed as possible from the program.

Three easy thoughts:

  1. This seems to allocate a new array:

Why not pass wẘ to α0 instead of w?

  1. Probk computes U0 or one of the Uk twice, depending on the value of k.

  2. Precompute the categorized a and pass it into Probk as an additional argument, so that agecat does not have to be called over and over again.

Then run the profiler and see where most of the time is actually spent.

1 Like

About your point 2, do you have any suggestions on how to avoid that?

What I meant was simply

function Probk(g,a,e,k,t; data=data, θ=θ)
    u00 = U0(g, a, e, t; data=data, θ=θ)
    ukk = Uk(g, a, e, k, t; data=data, θ=θ)
    uk = ( k == homecode ? u00 : ukk )
    1/(sum(exp(Uk(g, a, e, kk, t; data=data, θ=θ)-uk) for kk=1:5) + exp(u00-uk))
end
1 Like

Fantastic, thank you. I think there is another improvement, if it is possible to drop an element from 1:5 since for k != homecode, k is in 1:5, but that is precomputed and stored in ukk.

About point 1, doesn’t inlining take care of that?

1 Like

You could
sum(f(kk) for k = 1 : 5 if k != kk) + ukk

About point 1: do you mean a view? That would avoid the allocation, but is not entirely costless. I just don’t see the benefit of w[g, agecat(a), e, k, t] over wẘ[g, agecat(a), e, 1, k, t].

I don’t know if there is a reason to pass the entire wẘ array into U0 and Uk. If not, it would avoid a fair bit of indexing to just pass a view into a slice (over k).