From 68c0ace1961271b2b8586a294fd724e7c74dde0e Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Thu, 27 Aug 2026 13:49:32 -0400 Subject: [PATCH 1/3] Move connected-components helpers next to rescale! --- src/AbstractFixedEffectSolver.jl | 46 +++++++++++++++++++++++++++++- src/FixedEffect.jl | 48 -------------------------------- 2 files changed, 45 insertions(+), 49 deletions(-) diff --git a/src/AbstractFixedEffectSolver.jl b/src/AbstractFixedEffectSolver.jl index 430a5c2..523e7ec 100644 --- a/src/AbstractFixedEffectSolver.jl +++ b/src/AbstractFixedEffectSolver.jl @@ -201,7 +201,7 @@ 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). +# the first within each component. function normalize!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect}) idx = findall(fe -> isa(fe.interaction, UnitWeights), fes) length(idx) >= 2 && rescale!(view(fecoefs, idx), view(fes, idx)) @@ -231,3 +231,47 @@ function rescale!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVect end end end + +# Returns a vector of all connected 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 + +# 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 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 From abdde2f63bc399708d85d13bdf74b59126a83593 Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Thu, 27 Aug 2026 13:55:34 -0400 Subject: [PATCH 2/3] Replace Set-based BFS with union-find over group labels --- src/AbstractFixedEffectSolver.jl | 146 ++++++++++++++++++++----------- test/solve.jl | 33 +++++++ 2 files changed, 128 insertions(+), 51 deletions(-) diff --git a/src/AbstractFixedEffectSolver.jl b/src/AbstractFixedEffectSolver.jl index 523e7ec..d4f615f 100644 --- a/src/AbstractFixedEffectSolver.jl +++ b/src/AbstractFixedEffectSolver.jl @@ -209,69 +209,113 @@ function normalize!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVe 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 -# Returns a vector of all connected components -# A component is a vector that, for each fixed effect, -# contains all the refs that are included in the component. +# 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}) - 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 + 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 - push!(out, component_vec) + label[g] = c end end - return out + return labels, ncomponents end -# 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) +# 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 out + return 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 From 4f1f44cd0912e6334012b4c91f63cc6cd4cf0fd2 Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Thu, 27 Aug 2026 16:58:25 -0400 Subject: [PATCH 3/3] Inline fixed-effect coefficient normalization --- src/AbstractFixedEffectSolver.jl | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/src/AbstractFixedEffectSolver.jl b/src/AbstractFixedEffectSolver.jl index d4f615f..b278050 100644 --- a/src/AbstractFixedEffectSolver.jl +++ b/src/AbstractFixedEffectSolver.jl @@ -194,18 +194,13 @@ 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. -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})