From b5eef3ff4b39e49328ad8da8c642d4d646535471 Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Fri, 30 Oct 2026 14:02:36 -0400 Subject: [PATCH] Solve again on the residual when fixed effects explain most of a variable --- Project.toml | 5 +++-- src/AbstractFixedEffectSolver.jl | 19 ++++++++++++++++--- test/solve.jl | 16 ++++++++++++++++ 3 files changed, 35 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index b18d325..6285690 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FixedEffects" uuid = "c8885935-8500-56a7-9867-7708b20db0eb" -version = "3.5.0" +version = "3.5.1" [deps] GroupedArrays = "6407cd72-fade-4a84-8a1e-56e431fc1533" @@ -24,6 +24,7 @@ StatsBase = "0.33, 0.34" julia = "1.10" [extras] +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" @@ -32,4 +33,4 @@ PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["CategoricalArrays", "CUDA", "Metal", "Pkg", "PooledArrays", "Test"] +test = ["CategoricalArrays", "CUDA", "Metal", "Pkg", "PooledArrays", "Random", "Test"] diff --git a/src/AbstractFixedEffectSolver.jl b/src/AbstractFixedEffectSolver.jl index e94909f..98363dd 100644 --- a/src/AbstractFixedEffectSolver.jl +++ b/src/AbstractFixedEffectSolver.jl @@ -62,12 +62,25 @@ function solve_residuals!(r::AbstractVector{<:Real}, feM::AbstractFixedEffectSol iter, converged = 0, true if length(feM.m.plan.blocks) == 1 mul!(feM.x, feM.m', feM.b, 1, 0) + mul!(feM.r, feM.m, feM.x, -1, 1) else - _, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar, feM.g; atol = tol, btol = tol, maxiter = maxiter) - iter, converged = ch.mvps, ch.isconverged + # lsmr! stops relative to the norm of its right-hand side. When the fixed effects explain + # most of that norm (e.g. a variable with a large mean), this leaves an error that is large + # relative to the residual, so solve again on the residual. + size0 = sizeb = norm(feM.b) + while true + _, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar, feM.g; atol = tol, btol = tol, maxiter = maxiter - iter) + iter += ch.mvps + converged = ch.isconverged + mul!(feM.r, feM.m, feM.x, -1, 1) + sizer = norm(feM.r) + (converged && iter < maxiter && sizer < sqrt(tol) * sizeb && sizer > 100 * eps(T) * size0) || break + sizeb = sizer + copyto!(feM.b, feM.r) + fill!(feM.x, zero(T)) + end end converged || @warn "solve_residuals! did not converge within maxiter LSMR iterations; returned values may be inaccurate." iterations=iter maxiter tol - mul!(feM.r, feM.m, feM.x, -1, 1) if !(feM.weights isa UnitWeights) feM.r ./= sqrt.(feM.weights) end diff --git a/test/solve.jl b/test/solve.jl index ac70586..0f35ed3 100644 --- a/test/solve.jl +++ b/test/solve.jl @@ -1,3 +1,4 @@ +using Random using LinearAlgebra p1 = repeat(1:5, inner = 2) @@ -21,6 +22,21 @@ r_ols = [-0.2015993617092453, 0.2015993617092464, -0.2015993617092463, 0.2015 @test iterz == 0 @test all(iszero, rz) +@testset "residuals do not depend on the mean of the variable" begin + rng_mean = MersenneTwister(1) + n_mean = 2000 + fes_mean = [FixedEffect(rand(rng_mean, 1:100, n_mean)), FixedEffect(rand(rng_mean, 1:100, n_mean))] + x_mean = randn(rng_mean, n_mean) + r0 = solve_residuals!(copy(x_mean), fes_mean; tol = 1e-6)[1] + for m in (1e3, 1e6, 1e9) + r1 = solve_residuals!(x_mean .+ m, fes_mean; tol = 1e-6)[1] + @test norm(r1 - r0) < 2e-3 * norm(r0) + end + # a variable explained by the fixed effects + rfe = solve_residuals!(1e6 .+ fes_mean[1].refs .+ fes_mean[2].refs .* 1.0, fes_mean; tol = 1e-6)[1] + @test norm(rfe) < 1e-6 * sqrt(n_mean) +end + @testset "maxiter semantics" begin (r0, iter0, conv0) = @test_logs (:warn, r"solve_residuals!") solve_residuals!(deepcopy(x), fes; maxiter = 0) @test iter0 == 0