Mixed models: How to generate predictions "fixed" (marginalised?) over the random factors

[DominiqueMakowski]

Let's start with a simple linear mixed model that takes participant as random factors.

using Turing, StatsFuns, DataFrames, Random

# Generate correlated x and y data
data = DataFrame(x=1:100, y=1:100, participant=repeat(1:10, inner=10))
data.y = data.x .+ randn(100) .* 10

# Linear
@model function model_Linear(y, x, participant, n)

    # Priors - Fixed Effects
    μ_intercept ~ Normal(0, 0.5)
    μ_x ~ Normal(0, 0.5)

    σ ~ truncated(Normal(0.0, 1), lower=0)

    # Priors - Random Effects
    μ_intercept_random_sd ~ truncated(Normal(0.0, 0.3), lower=0)
    μ_x_random_sd ~ truncated(Normal(0, 0.3), lower=0)

    # Participant-level priors
    μ_intercept_random ~ filldist(Normal(0, μ_intercept_random_sd), n)
    μ_x_random ~ filldist(Normal(0, μ_x_random_sd), n)

    for i in 1:length(y)
        # Compute μ
        μ = μ_intercept + μ_intercept_random[participant[i]]
        μ += (μ_x + μ_x_random[participant[i]]) * x[i, 1]

        # Likelihood
        y[i] ~ Normal(μ, σ)
    end
end

posteriors = sample(model_Linear(data.y, data.x, data.participant, 10), NUTS(), 200)

I can easily generate predictions on the original data, useful for instance for posterior predictive checks.

# Predictions on whole dataset
pred = predict(model_Linear(fill(missing, 100), data.x, data.participant, 10), posteriors)

However, often for visualization purposes / postprocessing (marginal effects computation) we would need to generate predictions on new data, for instance on a data "grid" with specific values. While these values are often specifically chosen for the fixed effects, we are often interested in "marginalising" over the random effects, i.e., having the average prediction for the average random factor.

I thought that one way to go about it would be to also set participant as missing, but that doesn't seem to be the correct way.

# Predictions on new data "grid"
grid = DataFrame(x=[0, 50, 100], participant=[missing, missing, missing])

pred2 = predict(model_Linear(fill(missing, 3), grid.x, grid.participant, 3), posteriors)

Any advice on how to achieve something like that?

[torfjelde]

we are often interested in "marginalising" over the random effects, i.e., having the average prediction for the average random factor.

Maybe I'm misunderstanding, but do you just want

μ_intercept + μ_x * x[i, 1]

?

If you want to marginalize over the participants, then you need a notion of what the distribution of the participants are (which it doesn't seem like you have here, IIUC)

[DominiqueMakowski]

Thanks Tor as always

1. Vectorization issue

Unfortunately, it seems like when the model is in the vectorized form predicitons don't work (?)

@model function model_Linear(y, x, participant, n)

    # Priors - Fixed Effects
    μ_intercept ~ Normal(0, 0.5)
    μ_x ~ Normal(0, 0.5)

    σ ~ truncated(Normal(0.0, 1), lower=0)

    # Priors - Random Effects
    μ_intercept_random_sd ~ truncated(Normal(0.0, 0.3), lower=0)
    μ_x_random_sd ~ truncated(Normal(0, 0.3), lower=0)

    # Participant-level priors
    μ_intercept_random ~ filldist(Normal(0, μ_intercept_random_sd), n)
    μ_x_random ~ filldist(Normal(0, μ_x_random_sd), n)

    # Equation
    μ_fixed = μ_intercept .+ μ_x .* x[:, 1]
    μ_random = μ_intercept_random[participant] .+ μ_x_random[participant] .* x[:, 1]
    μ = μ_fixed .+ μ_random

    # Likelihood
    y .~ Normal.(μ, σ)
end

# Generate correlated x and y data
data = DataFrame(x=1:100, y=1:100, participant=repeat(1:10, inner=10))
data.y = data.x .+ randn(100) .* 10

# Sample
post = sample(model_Linear(data.y, data.x, data.participant, 10), NUTS(), 200)

# Predictions on whole dataset
pred = predict(model_Linear(fill(missing, 100), data.x, data.participant, 10), post)

ERROR: MethodError: no method matching loglikelihood(::Normal{Float64}, ::Missing)
The function `loglikelihood` exists, but no method is defined for this combination of argument types.

2. Using generated_quantities

You are correct I think, μ = μ_intercept + μ_intercept_random[participant[i]] + (μ_x + μ_x_random[participant[i]]) * x[i, 1] can be decomposed as:

μ_fixed = μ_intercept + μ_x * x[i, 1]
μ_random = μ_intercept_random[participant[i]] + μ_x_random[participant[i]] * x[i, 1]
μ = μ_fixed + μ_random

and thus μ_fixed is what I'm after. However, I am not sure how to access it during prediction on a new dataset (a new data grid of specific values of x). As far as I understand, generated_quantities() allows me to access the quantities generated during sampling?

3. Conditional inclusion if random part of equation

Finally, following your suggestion, I wrote the following model:

@model function model_Linear(y, x, participant, n)

    # Priors - Fixed Effects
    μ_intercept ~ Normal(0, 0.5)
    μ_x ~ Normal(0, 0.5)

    σ ~ truncated(Normal(0.0, 1), lower=0)

    # Priors - Random Effects
    μ_intercept_random_sd ~ truncated(Normal(0.0, 0.3), lower=0)
    μ_x_random_sd ~ truncated(Normal(0, 0.3), lower=0)

    # Participant-level priors
    μ_intercept_random ~ filldist(Normal(0, μ_intercept_random_sd), n)
    μ_x_random ~ filldist(Normal(0, μ_x_random_sd), n)

    for i in 1:length(y)
        # Compute μ
        μ = μ_intercept + μ_x * x[i, 1]
        if !ismissing(participant[i])
            μ += μ_intercept_random[participant[i]] + μ_x_random[participant[i]] * x[i, 1]
        end

        # Likelihood
        y[i] ~ Normal(μ, σ)
    end
end

While it works for original-data predictions:

# Generate correlated x and y data
data = DataFrame(x=1:100, y=1:100, participant=repeat(1:10, inner=10))
data.y = data.x .+ randn(100) .* 10

# Sample and predict
post = sample(model_Linear(data.y, data.x, data.participant, 10), NUTS(), 200)
pred = predict(model_Linear(fill(missing, 100), data.x, data.participant, 10), post)

It errors for predictions on a grid:

# Predictions on new data "grid"
grid = DataFrame(x=[0, 50, 100], participant=[missing, missing, missing])
pred2 = predict(model_Linear(fill(missing, 3), grid.x, grid.participant, 3), post)

ERROR: BoundsError: attempt to access 3-element Vector{Float64} at index [4]

I feel like I'm close, but the I can't make sense of the error 🤔

[torfjelde]

As far as I understand, generated_quantities() allows me to access the quantities generated during sampling?

So generated_quantities extracts the return statement of the @model conditioned on the parameter realizations in chain. And this is exactly what you want:)

However one issue with this, is that in your chain you'll have variables μ_intercept_random that are of length n (4, in your case), which might be different for your prediction popluation (3, in your case). This is why you're encountering the BoundsError: predict tries to find a place to set the variable μ_intercept_random[4] (which was present during sampling), but now μ_intercept_random is only of length 3 (since you set n=3).

So the way to address this is to find a way to make your model not have the participant-level random variables involved when you're predicting (since that will always fail if you're n is incorrect). A simple way to do that is

@model function model_Linear(y, x, participant, n)

    # Priors - Fixed Effects
    μ_intercept ~ Normal(0, 0.5)
    μ_x ~ Normal(0, 0.5)

    σ ~ truncated(Normal(0.0, 1), lower=0)

    # Priors - Random Effects
    μ_intercept_random_sd ~ truncated(Normal(0.0, 0.3), lower=0)
    μ_x_random_sd ~ truncated(Normal(0, 0.3), lower=0)

    # Participant-level priors (we only include these if `participant` is not `nothing`)
    if participant !== nothing
        μ_intercept_random ~ filldist(Normal(0, μ_intercept_random_sd), n)
        μ_x_random ~ filldist(Normal(0, μ_x_random_sd), n)
    end

    for i in 1:length(y)
        # Compute μ
        μ = μ_intercept + μ_x * x[i, 1]
        if participant !== nothing
            μ += μ_intercept_random[participant[i]] + μ_x_random[participant[i]] * x[i, 1]
        end

        # Likelihood
        y[i] ~ Normal(μ, σ)
    end

    # For `generated_quantities`
    return (μ_fixed = μ_intercept .+ μ_x .* x[:, 1],)
end

# Convenience method for when participants aren't available.
model_Linear(y, x) = model_Linear(y, x, nothing, nothing)

Then you can do

# Predictions on new data "grid"
grid = DataFrame(x=[0, 50, 100], participant=[missing, missing, missing])
model_predict = model_Linear(fill(missing, 3), grid.x)
pred2 = predict(model_predict, post)
results = generated_quantities(model_predict, post)

[torfjelde]

If you really want to make your life a bit nicer, you could also separate the fixed and random effects into two different models:

"""
    model_linear_base(x)

Model component of linear model with mixed effects.
"""
@model function model_linear_base(x)
    # Priors - Fixed Effects
    μ_intercept ~ Normal(0, 0.5)
    μ_x ~ Normal(0, 0.5)

    # Priors - Random Effects
    μ_intercept_random_sd ~ truncated(Normal(0.0, 0.3), lower=0)
    μ_x_random_sd ~ truncated(Normal(0, 0.3), lower=0)

    return μ_intercept .+ μ_x .* x[:, 1]
end

"""
    model_linear_participant(x, participant, n)

Model component of linear model with participant-specific effects.
"""
@model function model_linear_participant(x, participant, n)
    μ_intercept_random ~ filldist(Normal(0, μ_intercept_random_sd), n)
    μ_x_random ~ filldist(Normal(0, μ_x_random_sd), n)

    return μ_intercept_random[participant] .+ μ_x_random[participant] .* x[:, 1]
end

@model function model_linear(y, x, participant, n)
    @submodel μ_base = model_linear_base(x)
    μ = μ_base

    if participant !== nothing
        @submodel μ_participant = model_linear_participant(x, participant, n)
        μ = μ + μ_participant
    end

    # Observe.
    σ ~ truncated(Normal(0.0, 1), lower=0)
    y .~ Normal.(μ, σ)
end

[DominiqueMakowski]

Thanks that's nice! I somehow got invested in Turing's development though, and it makes me wonder if the @submodel syntax is future-proof (cf. discussions around returned_quantities etc)

[torfjelde]

and it makes me wonder if the @submodel syntax is future-proof (cf. discussions around returned_quantities etc)

@submodel would not just be removed out of nowhere:) It'll still be around for quite some time, no doubt. But yes, we're thinking about renaming it to @returned_quantities or something similar, though the underlying functionality will be identical with very minor syntax changes (specifically, @submodel x = model becomes x = @submodel model).

So wouldnt' say this should be cause from refraining to use it:)

[DominiqueMakowski]

Super sorry to keep annoying with this Tor, but now using the vectorized syntax suggested above it looks like the predict() methods got broken...

Here's an MWE:

using Turing, StatsFuns, DataFrames, Random


@model function mod(y, x1, x2)

    μ_intercept ~ Normal(0, 0.5)
    μ_x1 ~ Normal(1, 0.5)
    μ_x2 ~ Normal(2, 0.5)
    σ ~ truncated(Normal(0.0, 1), lower=0)

    μ = μ_intercept .+ μ_x1 .* x1 .+ μ_x2 .* x2
    y .~ Normal.(μ, σ)
end


fit = mod(rand(10), rand(10), rand(10))
chain = sample(fit, NUTS(), 400)
pred = predict(mod([missing for i in 1:10], rand(10), rand(10)), chain)

ERROR: MethodError: no method matching loglikelihood(::Normal{Float64}, ::Missing)
The function `loglikelihood` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  loglikelihood(::Distribution{ArrayLikeVariate{N}}, ::AbstractArray{<:Real, M}) where {N, M}
   @ Distributions C:\Users\domma\.julia\packages\Distributions\uuqsE\src\common.jl:444
  loglikelihood(::Distribution{ArrayLikeVariate{N}}, ::AbstractArray{<:AbstractArray{<:Real, N}}) where N
   @ Distributions C:\Users\domma\.julia\packages\Distributions\uuqsE\src\common.jl:461
  loglikelihood(::Model, ::Any)
   @ DynamicPPL C:\Users\domma\.julia\packages\DynamicPPL\aKhfc\src\simple_varinfo.jl:662
  ...

My guess is that we should in this example use MvNormal, but is there a way to make it work for likelihood distributions that don't have a multivariate form?

[torfjelde]

Ah sorry, that's my bad. .~ for variables passed as arguments doesn't for like this.

For this, you'll have to use condition / model | (y = ...,), i.e.

@model function mod(x1, x2, ::Type{TV}=Vector{Float64}) where {TV}

    μ_intercept ~ Normal(0, 0.5)
    μ_x1 ~ Normal(1, 0.5)
    μ_x2 ~ Normal(2, 0.5)
    σ ~ truncated(Normal(0.0, 1), lower=0)

    μ = μ_intercept .+ μ_x1 .* x1 .+ μ_x2 .* x2
    y = TV(undef, length(μ))
    y .~ Normal.(μ, σ)
end


fit = mod(rand(10), rand(10)) | (y = rand(10),)
chain = sample(fit, NUTS(), 400)
pred = predict(decondition(fit, :y), chain)

This is because we check the left-hand side of .~ to determine if it's conditioned or not. In the case where you pass y as an argument, whether or not it's conditioned is determined by whether or not it equals missing. In your case, the left-hand side, i.e. y, does not equal missing (but is instead a vector of missings), and so is treated as an observation instead of random.

We could implement it so that it checks each entry of the left-hand side, i.e. y in this case, but this will then be equivalent to the for-loop we wrote earlier, in which case the .~ becomes sort of useless 🤷