Fix singleton pruning across fixed effects

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"

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
 
 

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)