Linear model with categorical variable

General Usage
2

This allocates an array μall with eltype Float64, but μ1 and μ2 have eltype Dual because of automatic differentiation with ForwardDiff. This would work:

    μall = zeros(size(S), Base.promote_eltype(μ1, μ2))

A better approach would be to compute S .== 0 and S .== 1 outside of the model evaluation. e.g. do

@model function linear_model(X, Y, S, S0_inds = findall(iszero, S), S1_inds = findall(isone, S), S_perm = invperm([S0_inds; S1_inds]))
    ...
    μ1 = @views α1 .+ β .* (X[S1_inds] .- mean(X[S1_inds]));
    μ2 = @views α2 .+ β .* (X[S0_inds] .- mean(X[S0_inds]));
    μall = vcat(μ2, μ1)[S_perm]
    ...

There are probably even simpler ways to do this. e.g. you could center your X values outside the model, since that’s wasted computation for each gradient evaluation. And you should check that this is equivalent to your model.

1 Like