From f0fe3eaf9038896230f8555975cee8694424fa62 Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Tue, 25 Aug 2026 10:01:30 -0400 Subject: [PATCH 1/4] Fix singleton pruning across fixed effects The previous repeated-scan loop could stop after a clean pass on the first fixed effect and miss singleton levels defined by a later fixed effect. That left observations in the sample when their singleton status came from the second or subsequent FE dimension. Replace repeated scans with a queue-based peel over FE group counts, using active-index xors to identify singleton observations without building reverse offsets. Add coverage for later-FE singletons, cascades, inactive rows, interaction FEs, duplicate singleton queue entries, and partial_out compatibility. Bump version to 1.13.3. --- Project.toml | 2 +- src/utils/fixedeffects.jl | 104 ++++++++++++++++++++++++++++---------- test/fit.jl | 62 +++++++++++++++++++++++ 3 files changed, 139 insertions(+), 29 deletions(-) diff --git a/Project.toml b/Project.toml index 3f95c06..d70c646 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FixedEffectModels" uuid = "9d5cd8c9-2029-5cab-9928-427838db53e3" -version = "1.13.2" +version = "1.13.3" [deps] DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" diff --git a/src/utils/fixedeffects.jl b/src/utils/fixedeffects.jl index e653128..7e70a3f 100644 --- a/src/utils/fixedeffects.jl +++ b/src/utils/fixedeffects.jl @@ -4,43 +4,91 @@ ## ############################################################################## -function drop_singletons!(esample, fes::Vector{<:FixedEffect}) +function drop_singletons!(esample, fes::AbstractVector{<:FixedEffect}, _nthreads::Integer = Threads.nthreads()) + isempty(fes) && return 0 + + counts = Vector{Vector{Int}}(undef, length(fes)) + # When a group count is one, the xor of active observation indices identifies + # the only remaining observation in that group. + active_xor = Vector{Vector{Int}}(undef, length(fes)) + queue_fe = Int[] + queue_ref = Int[] + + # Count active observations by FE level and seed the singleton queue. + @inbounds for j in eachindex(fes) + fe = fes[j] + counts[j] = zeros(Int, fe.n) + active_xor[j] = zeros(Int, fe.n) + counts_j = counts[j] + nactive = count_groups!(counts_j, active_xor[j], esample, fe.refs) + nactive == 0 && return 0 + + nsingleton_groups = 0 + @inbounds for ref in eachindex(counts_j) + if counts_j[ref] == 1 + nsingleton_groups += 1 + end + end + if nsingleton_groups == nactive + fill!(esample, false) + return nactive + end + if nsingleton_groups > 0 + @inbounds for ref in eachindex(counts_j) + if counts_j[ref] == 1 + push!(queue_fe, j) + push!(queue_ref, ref) + end + end + end + end + isempty(queue_fe) && return 0 + nsingletons = 0 - ncleanpasses = 0 - caches = [Vector{UInt8}(undef, fes[i].n) for i in eachindex(fes)] - for (fe, cache) in Iterators.cycle(zip(fes,caches)) - n = drop_singletons!(esample, fe, cache) - nsingletons += n - if n > 0 - # found singletons, reset the counter - ncleanpasses = 0 - else - # otherwise, increment counter - ncleanpasses += 1 + head = 1 + # Peel singleton FE levels until every remaining level has degree at least two. + @inbounds while head <= length(queue_fe) + j = queue_fe[head] + ref = queue_ref[head] + head += 1 + counts[j][ref] == 1 || continue + + obsindex = active_xor[j][ref] + obsindex != 0 && esample[obsindex] || continue + + esample[obsindex] = false + nsingletons += 1 + # Removing one observation can create new singleton levels in every FE. + for k in eachindex(fes) + ref_k = fes[k].refs[obsindex] + counts_k = counts[k] + active_xor_k = active_xor[k] + counts_k[ref_k] -= 1 + active_xor_k[ref_k] = xor(active_xor_k[ref_k], obsindex) + if counts_k[ref_k] == 1 + push!(queue_fe, k) + push!(queue_ref, ref_k) + end end - # if the last N-1 passes have not found singletons (where N is number of FE groups), break the loop - ncleanpasses >= length(fes) - 1 && break end return nsingletons end -function drop_singletons!(esample, fe::FixedEffect, cache) - refs = fe.refs - fill!(cache, 0) - @inbounds for i in eachindex(esample, refs) # count obs in each FE group - if esample[i] - # no need to keep counting obs after 2 (counters are 8-bit integers) - cache[refs[i]] = min(0x02, cache[refs[i]] + 0x01) - end - end - n = 0 +function count_groups!(counts, active_xor, esample, refs) + nactive = 0 @inbounds for i in eachindex(esample, refs) - if esample[i] && cache[refs[i]] == 0x01 - esample[i] = false - n += 1 + if esample[i] + nactive += 1 + ref = refs[i] + counts[ref] += 1 + active_xor[ref] = xor(active_xor[ref], i) end end - return n + return nactive +end + +function drop_singletons!(esample, fe::FixedEffect, cache) + return drop_singletons!(esample, FixedEffect[fe]) end diff --git a/test/fit.jl b/test/fit.jl index fe88a30..fbd9ea6 100644 --- a/test/fit.jl +++ b/test/fit.jl @@ -604,6 +604,68 @@ end @testset "singletons" begin + @testset "singleton pruning k-core" begin + esample = trues(7) + fes = FixedEffectModels.FixedEffect[ + FixedEffectModels.FixedEffect([1, 1, 2, 2, 3, 3, 4]), + FixedEffectModels.FixedEffect([1, 2, 1, 2, 2, 3, 3]), + ] + @test FixedEffectModels.drop_singletons!(esample, fes) == 3 + @test esample == Bool[true, true, true, true, false, false, false] + + esample = Bool[true, true, true, true, false] + fes = FixedEffectModels.FixedEffect[ + FixedEffectModels.FixedEffect([1, 1, 2, 2, 3]), + FixedEffectModels.FixedEffect([1, 2, 1, 2, 3]), + ] + @test FixedEffectModels.drop_singletons!(esample, fes) == 0 + @test esample == Bool[true, true, true, true, false] + + esample = trues(5) + fes = FixedEffectModels.FixedEffect[ + FixedEffectModels.FixedEffect([1, 1, 1, 2, 2]), + FixedEffectModels.FixedEffect([1, 2, 3, 2, 3]), + ] + @test FixedEffectModels.drop_singletons!(esample, fes) == 1 + @test esample == Bool[false, true, true, true, true] + + esample = Bool[true, false, true, true] + fes = FixedEffectModels.FixedEffect[ + FixedEffectModels.FixedEffect([1, 1, 2, 2]), + ] + @test FixedEffectModels.drop_singletons!(esample, fes) == 1 + @test esample == Bool[false, false, true, true] + + esample = trues(4) + fes = FixedEffectModels.FixedEffect[ + FixedEffectModels.FixedEffect([1, 1, 2, 2], [1, 2, 1, 2]), + ] + @test FixedEffectModels.drop_singletons!(esample, fes) == 4 + @test esample == falses(4) + + esample = trues(1) + fes = FixedEffectModels.FixedEffect[ + FixedEffectModels.FixedEffect([1]), + FixedEffectModels.FixedEffect([1]), + ] + @test FixedEffectModels.drop_singletons!(esample, fes) == 1 + @test esample == Bool[false] + + esample = trues(3) + fes = FixedEffectModels.FixedEffect[ + FixedEffectModels.FixedEffect([1, 1, 2]; interaction = [0.0, 0.0, 0.0]), + ] + @test FixedEffectModels.drop_singletons!(esample, fes) == 1 + @test esample == Bool[true, true, false] + + df_partial = DataFrame(y = [1.0, 2.0, 3.0], g = [1, 1, 2]) + _, esample_partial, _, _, _ = partial_out(df_partial, @formula(y ~ fe(g)); drop_singletons = true) + @test esample_partial == Bool[true, true, false] + + esample = Bool[true, false, true] + @test FixedEffectModels.drop_singletons!(esample, FixedEffectModels.FixedEffect[]) == 0 + @test esample == Bool[true, false, true] + end df = DataFrame(CSV.File(joinpath(dirname(pathof(FixedEffectModels)), "../dataset/Cigar.csv"))) df.StateC = categorical(df.State) From 4159d73f367c65ffa2013f233583bba7b4064c7a Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Tue, 25 Aug 2026 10:11:48 -0400 Subject: [PATCH 2/4] Remove brittle exact iteration assertion Solver iteration reporting can vary across dependency versions when the coefficients are already computed without an iterative solve. Keep the coefficient assertion and remove the exact x.iterations == 1 check. --- test/fit.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/fit.jl b/test/fit.jl index fbd9ea6..5c89645 100644 --- a/test/fit.jl +++ b/test/fit.jl @@ -31,7 +31,6 @@ using CUDA, Metal m = @formula Sales ~ Price + fe(State) x = reg(df, m) @test coef(x) ≈ [-0.20984] atol = 1e-4 - @test x.iterations == 1 m = @formula Sales ~ Price + fe(State) + fe(Year) x = reg(df, m) From d9789721ce4e0f95e98f26e4cdd699b24a3fbe4a Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Tue, 25 Aug 2026 15:04:27 -0400 Subject: [PATCH 3/4] Clarify singleton pruning strategy --- src/utils/fixedeffects.jl | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/src/utils/fixedeffects.jl b/src/utils/fixedeffects.jl index 7e70a3f..e531d39 100644 --- a/src/utils/fixedeffects.jl +++ b/src/utils/fixedeffects.jl @@ -4,9 +4,19 @@ ## ############################################################################## +# Strategy: use a k-core peeling algorithm, with k = 2. Treat observations and +# fixed-effect levels as an incidence graph (a graph connecting each observation +# to the FE levels it belongs to) and repeatedly peel singleton FE levels until +# every remaining FE level contains at least two active observations (the graph's +# 2-core). +# For each FE level, keep its active-observation count and the xor of active +# observation indices. When the count is one, the xor identifies the single +# observation to drop. Dropping an observation updates all FE levels it belongs +# to and queues any levels that become singletons. function drop_singletons!(esample, fes::AbstractVector{<:FixedEffect}, _nthreads::Integer = Threads.nthreads()) isempty(fes) && return 0 + # counts[j][ref] is the number of active observations in FE j with level ref. counts = Vector{Vector{Int}}(undef, length(fes)) # When a group count is one, the xor of active observation indices identifies # the only remaining observation in that group. @@ -17,12 +27,15 @@ function drop_singletons!(esample, fes::AbstractVector{<:FixedEffect}, _nthreads # Count active observations by FE level and seed the singleton queue. @inbounds for j in eachindex(fes) fe = fes[j] - counts[j] = zeros(Int, fe.n) - active_xor[j] = zeros(Int, fe.n) - counts_j = counts[j] - nactive = count_groups!(counts_j, active_xor[j], esample, fe.refs) + counts_j = zeros(Int, fe.n) + active_xor_j = zeros(Int, fe.n) + # nactive is the number of observations currently kept in the estimation + # sample, i.e. the number of true entries in esample. + # count_groups! also fills counts_j and active_xor_j. + nactive = count_groups!(counts_j, active_xor_j, esample, fe.refs) nactive == 0 && return 0 + # Count FE levels that currently contain exactly one active observation. nsingleton_groups = 0 @inbounds for ref in eachindex(counts_j) if counts_j[ref] == 1 @@ -41,6 +54,8 @@ function drop_singletons!(esample, fes::AbstractVector{<:FixedEffect}, _nthreads end end end + counts[j] = counts_j + active_xor[j] = active_xor_j end isempty(queue_fe) && return 0 From 21b0b2347137daaf2c93d68ca41dc64c6b2d2d80 Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Tue, 25 Aug 2026 15:08:12 -0400 Subject: [PATCH 4/4] Use nobs in singleton pruning --- src/utils/fixedeffects.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/utils/fixedeffects.jl b/src/utils/fixedeffects.jl index e531d39..5080096 100644 --- a/src/utils/fixedeffects.jl +++ b/src/utils/fixedeffects.jl @@ -29,11 +29,11 @@ function drop_singletons!(esample, fes::AbstractVector{<:FixedEffect}, _nthreads fe = fes[j] counts_j = zeros(Int, fe.n) active_xor_j = zeros(Int, fe.n) - # nactive is the number of observations currently kept in the estimation + # nobs is the number of observations currently kept in the estimation # sample, i.e. the number of true entries in esample. # count_groups! also fills counts_j and active_xor_j. - nactive = count_groups!(counts_j, active_xor_j, esample, fe.refs) - nactive == 0 && return 0 + nobs = count_groups!(counts_j, active_xor_j, esample, fe.refs) + nobs == 0 && return 0 # Count FE levels that currently contain exactly one active observation. nsingleton_groups = 0 @@ -42,9 +42,9 @@ function drop_singletons!(esample, fes::AbstractVector{<:FixedEffect}, _nthreads nsingleton_groups += 1 end end - if nsingleton_groups == nactive + if nsingleton_groups == nobs fill!(esample, false) - return nactive + return nobs end if nsingleton_groups > 0 @inbounds for ref in eachindex(counts_j) @@ -90,16 +90,16 @@ function drop_singletons!(esample, fes::AbstractVector{<:FixedEffect}, _nthreads end function count_groups!(counts, active_xor, esample, refs) - nactive = 0 + nobs = 0 @inbounds for i in eachindex(esample, refs) if esample[i] - nactive += 1 + nobs += 1 ref = refs[i] counts[ref] += 1 active_xor[ref] = xor(active_xor[ref], i) end end - return nactive + return nobs end function drop_singletons!(esample, fe::FixedEffect, cache)