Speed up connected-components normalization

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

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
 

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