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..5080096 100644 --- a/src/utils/fixedeffects.jl +++ b/src/utils/fixedeffects.jl @@ -4,43 +4,106 @@ ## ############################################################################## -function drop_singletons!(esample, fes::Vector{<:FixedEffect}) +# 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. + 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) + # 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. + 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 + @inbounds for ref in eachindex(counts_j) + if counts_j[ref] == 1 + nsingleton_groups += 1 + end + end + if nsingleton_groups == nobs + fill!(esample, false) + return nobs + 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 + counts[j] = counts_j + active_xor[j] = active_xor_j + 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) + nobs = 0 @inbounds for i in eachindex(esample, refs) - if esample[i] && cache[refs[i]] == 0x01 - esample[i] = false - n += 1 + if esample[i] + nobs += 1 + ref = refs[i] + counts[ref] += 1 + active_xor[ref] = xor(active_xor[ref], i) end end - return n + return nobs +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..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) @@ -604,6 +603,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)