diff --git a/src/partial_out.jl b/src/partial_out.jl index 30f866c..9e94bfc 100644 --- a/src/partial_out.jl +++ b/src/partial_out.jl @@ -10,6 +10,7 @@ Partial out variables in a Dataframe * `double_precision::Bool`: Should the demeaning operation use Float64 rather than Float32? Default to true. * `tol::Real`: Tolerance * `align::Bool`: Should the returned DataFrame align with the original DataFrame in case of missing values? Default to true. +* `drop_singletons::Bool=false`: Should singletons be dropped? ### Returns * `::DataFrame`: a dataframe with as many columns as there are dependent variables and as many rows as the original dataframe. @@ -40,7 +41,8 @@ function partial_out( @nospecialize(method::Symbol = :cpu), @nospecialize(double_precision::Bool = true), @nospecialize(tol::Real = double_precision ? 1e-8 : 1e-6), - @nospecialize(align = true)) + @nospecialize(align = true), + @nospecialize(drop_singletons = false)) if (ConstantTerm(0) ∉ eachterm(f.rhs)) & (ConstantTerm(1) ∉ eachterm(f.rhs)) f = FormulaTerm(f.lhs, tuple(ConstantTerm(1), eachterm(f.rhs)...)) @@ -67,6 +69,7 @@ function partial_out( fes, ids, ids_fes = parse_fixedeffect(df, formula_fes) has_fes = !isempty(fes) + drop_singletons && drop_singletons!(esample, fes, Threads.nthreads()) nobs = sum(esample) (nobs > 0) || throw("sample is empty") @@ -135,6 +138,8 @@ function partial_out( residuals .= residuals .+ m end + dof_fes = mapreduce(nunique, +, fes, init=0) + # Return a dataframe out = DataFrame() j = 0 @@ -148,5 +153,5 @@ function partial_out( out[!, Symbol(y)] = residuals[:, j] end end - return out, esample, iterations, convergeds + return out, esample, iterations, convergeds, dof_fes end