diff --git a/src/AbstractFixedEffectSolver.jl b/src/AbstractFixedEffectSolver.jl index 430a5c2..b278050 100644 --- a/src/AbstractFixedEffectSolver.jl +++ b/src/AbstractFixedEffectSolver.jl @@ -194,40 +194,123 @@ function recover_coefficients(::Type{T}, fes::Vector{<:FixedEffect}, plan::Absor end end end - normalize!(group_coefs, fes) - return Vector{Tout}[Tout.(coef[fe.refs]) for (coef, fe) in zip(group_coefs, fes)] -end - -# Fixed-effect coefficients are generally not unique: within each connected -# component, a constant can be shifted between the scalar (non-interacted) -# fixed effects. Pin down a solution by demeaning every scalar fixed effect but -# the first within each component (uses `components` from FixedEffect.jl). -function normalize!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect}) + # Fixed-effect coefficients are generally not unique: within each connected + # component, a constant can be shifted between the scalar (non-interacted) + # fixed effects. Pin down a solution by demeaning every scalar fixed effect but + # the first within each component. idx = findall(fe -> isa(fe.interaction, UnitWeights), fes) - length(idx) >= 2 && rescale!(view(fecoefs, idx), view(fes, idx)) - return fecoefs + length(idx) >= 2 && rescale!(view(group_coefs, idx), view(fes, idx)) + return Vector{Tout}[Tout.(coef[fe.refs]) for (coef, fe) in zip(group_coefs, fes)] end function rescale!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect}) - for component_vec in components(fes) - m = 0.0 - # demean all fixed effects except the first - for j in length(fecoefs):(-1):2 - fecoef, component = fecoefs[j], component_vec[j] - mj = 0.0 - for k in component - mj += fecoef[k] + labels, ncomponents = components(fes) + shift = zeros(ncomponents) # per component, total mean moved to the first fixed effect + sums = zeros(ncomponents) + counts = zeros(Int, ncomponents) + # demean all fixed effects except the first + for j in length(fecoefs):(-1):2 + fecoef, label = fecoefs[j], labels[j] + fill!(sums, 0.0) + fill!(counts, 0) + for g in eachindex(label) + c = label[g] + if c > 0 + sums[c] += fecoef[g] + counts[c] += 1 end - mj = mj / length(component) - for k in component - fecoef[k] -= mj + end + for g in eachindex(label) + c = label[g] + if c > 0 + fecoef[g] -= sums[c] / counts[c] + end + end + for c in 1:ncomponents + if counts[c] > 0 + shift[c] += sums[c] / counts[c] end - m += mj end - # rescale the first fixed effects - fecoef, component = fecoefs[1], component_vec[1] - for k in component - fecoef[k] += m + end + # rescale the first fixed effect + fecoef, label = fecoefs[1], labels[1] + for g in eachindex(label) + c = label[g] + if c > 0 + fecoef[g] += shift[c] + end + end +end + +# Connected components of the graph linking groups of different fixed effects +# through shared observations, via union-find over the group labels of all +# fixed effects. Returns, for each fixed effect, a vector mapping each group to +# its component id (0 for a group with no observation, which can arise from +# subsetting), and the number of components. +function components(fes::AbstractVector{<:FixedEffect}) + offsets = Vector{Int}(undef, length(fes) + 1) + offsets[1] = 0 + for (j, fe) in enumerate(fes) + offsets[j + 1] = offsets[j] + fe.n + end + # each observation links its first-effect group to its group in every other effect + parent = collect(1:offsets[end]) + treesize = ones(Int, offsets[end]) + refs1 = fes[1].refs + for j in 2:length(fes) + refsj = fes[j].refs + offset = offsets[j] + for i in eachindex(refs1) + _union!(parent, treesize, Int(refs1[i]), offset + Int(refsj[i])) end end + seen = falses(offsets[end]) + for (j, fe) in enumerate(fes) + offset = offsets[j] + for r in fe.refs + seen[offset + Int(r)] = true + end + end + labels = Vector{Int}[zeros(Int, fe.n) for fe in fes] + component_of_root = zeros(Int, offsets[end]) + ncomponents = 0 + for (j, fe) in enumerate(fes) + label, offset = labels[j], offsets[j] + for g in 1:fe.n + seen[offset + g] || continue + root = _find!(parent, offset + g) + c = component_of_root[root] + if c == 0 + ncomponents += 1 + c = ncomponents + component_of_root[root] = c + end + label[g] = c + end + end + return labels, ncomponents +end + +# find with path halving +function _find!(parent::Vector{Int}, i::Int) + @inbounds while parent[i] != i + parent[i] = parent[parent[i]] + i = parent[i] + end + return i +end + +# union by size +function _union!(parent::Vector{Int}, treesize::Vector{Int}, i::Int, j::Int) + ri = _find!(parent, i) + rj = _find!(parent, j) + ri == rj && return + if treesize[ri] < treesize[rj] + ri, rj = rj, ri + end + @inbounds begin + parent[rj] = ri + treesize[ri] += treesize[rj] + end + return end diff --git a/src/FixedEffect.jl b/src/FixedEffect.jl index 70f5e54..2564ae2 100644 --- a/src/FixedEffect.jl +++ b/src/FixedEffect.jl @@ -54,52 +54,4 @@ Base.getindex(fe::FixedEffect, ::Colon) = fe @inbounds interaction = fe.interaction[esample] return FixedEffect{typeof(fe.refs), typeof(fe.interaction)}(refs, interaction, fe.n) end -############################################################################## -## -## Find connected components -## -############################################################################## -# Return a vector of sets that contains the indices of each unique value -function refsrev(fe::FixedEffect) - out = Vector{Int}[Int[] for _ in 1:fe.n] - for i in eachindex(fe.refs) - push!(out[fe.refs[i]], i) - end - return out -end - -# Returns a vector of all components -# A component is a vector that, for each fixed effect, -# contains all the refs that are included in the component. -function components(fes::AbstractVector{<:FixedEffect}) - refs_vec = Vector{Int}[fe.refs for fe in fes] - refsrev_vec = Vector{Vector{Int}}[refsrev(fe) for fe in fes] - visited = falses(length(refs_vec[1])) - out = Vector{Set{Int}}[] - for i in eachindex(visited) - if !visited[i] - # obs not visited yet, so create new component - component_vec = Set{Int}[Set{Int}() for _ in 1:length(refsrev_vec)] - # visit all obs in the same components - tovisit = Set{Int}(i) - while !isempty(tovisit) - for (component, refs, refsrev) in zip(component_vec, refs_vec, refsrev_vec) - ref = refs[i] - # if group is not in component yet - if ref ∉ component - # add group to the component - push!(component, ref) - # visit other observations in same group - union!(tovisit, refsrev[ref]) - end - end - # mark obs as visited - i = pop!(tovisit) - visited[i] = true - end - push!(out, component_vec) - end - end - return out -end diff --git a/test/solve.jl b/test/solve.jl index 9c3e505..0370a3d 100644 --- a/test/solve.jl +++ b/test/solve.jl @@ -300,3 +300,36 @@ end FixedEffects._USE_THREADED_GATHER[] = true @test r_default ≈ r_serial atol = 1e-8 end + +@testset "connected components normalization" begin + # two components: {p1 ∈ (1, 2)} × {p2 ∈ (1, 2)} and {p1 ∈ (3, 4)} × {p2 = 3} + p1c = [1, 1, 2, 2, 3, 3, 4] + p2c = [1, 2, 1, 2, 3, 3, 3] + yc = [0.1, 1.3, -0.4, 0.7, 2.0, -1.1, 0.6] + fes_c = [FixedEffect(p1c), FixedEffect(p2c)] + labels, ncomponents = FixedEffects.components(fes_c) + @test ncomponents == 2 + @test labels[1][1] == labels[1][2] == labels[2][1] == labels[2][2] + @test labels[1][3] == labels[1][4] == labels[2][3] + @test labels[1][1] != labels[1][3] + coefs_c, _, conv_c = solve_coefficients!(copy(yc), fes_c) + @test conv_c + rc = solve_residuals!(copy(yc), fes_c)[1] + @test yc .- coefs_c[1] .- coefs_c[2] ≈ rc atol = 1e-6 + # the second fixed effect has mean zero over the groups of each component + @test coefs_c[2][1] + coefs_c[2][2] ≈ 0 atol = 1e-6 + @test coefs_c[2][5] ≈ 0 atol = 1e-6 + + # groups with no observation (from subsetting) are ignored + fes_s = [FixedEffect(p1c)[1:4], FixedEffect(p2c)[1:4]] + labels_s, ncomponents_s = FixedEffects.components(fes_s) + @test ncomponents_s == 1 + @test labels_s[1][3:4] == [0, 0] + @test labels_s[2][3] == 0 + ys = yc[1:4] + coefs_s, _, conv_s = solve_coefficients!(copy(ys), fes_s) + @test conv_s + @test all(all(isfinite, coef) for coef in coefs_s) + rs = solve_residuals!(copy(ys), fes_s)[1] + @test ys .- coefs_s[1] .- coefs_s[2] ≈ rs atol = 1e-6 +end