JuliaStats / GLM.jl

Generalized linear models in Julia
Other
584 stars 114 forks source link

type stability of vector of distributions #561

Open biona001 opened 1 month ago

biona001 commented 1 month ago

I have a vector y which is a Vector{Float64}, but each y[i] is sampled from a different univariate distribution specified by vecdist[i]. I want to compute some loglikelihood function.

I was surprised functions like GLM.linkinv and GLM.loglik_obs are not type stable? Is this expected behavior?

MWE

using Distributions
using GLM
using Random
import GLM.loglik_obs, GLM.linkinv

function component_loglikelihood(y, vecdist, veclink, η)
    logl = zero(eltype(y))
    for j in eachindex(y)
        dist = vecdist[j]
        link = veclink[j]
        μ_j = GLM.linkinv(link, η[j])
        logl += GLM.loglik_obs(dist, y[j], μ_j, 1.0, 1.0)
    end
    return logl
end

# simulate data
n = 10
vecdist = rand([Bernoulli(), Poisson(), Normal()], n)
veclink = canonicallink.(vecdist)
y = [rand(dist) for dist in vecdist] |> Vector{Float64}
η = randn(n)

# check type
@code_warntype component_loglikelihood(y, vecdist, veclink, η)

The output is shown below (apologize for the picture, copy-pasting into github won't highlight the red colors)

Screenshot 2024-05-17 at 2 09 33 PM

From the picture, logl::Any even though I initialized it as zero(eltype(y)), and μ_j::Any even though μ_j=GLM.linkinv(link, η[j]).