slightly better performance (#257)

* Update fit.jl

* Update fit.jl

* Update fixedeffects.jl

* Update fit.jl

* Update fixedeffects.jl

* Update fixedeffects.jl

* Update fixedeffects.jl

* Update fixedeffects.jl

* Update fixedeffects.jl

* simplify parsing

* update

* simplify a bit

* Update fit.jl

* Update fit.jl

* Update fit.jl

* Update fit.jl

* remove dof_ad

* Update fit.jl

Author: matthieugomez

Project.toml

@@ -1,6 +1,6 @@
 name = "FixedEffectModels"
 uuid = "9d5cd8c9-2029-5cab-9928-427838db53e3"
-version = "1.10.2"
+version = "1.11.0"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"

src/FixedEffectModel.jl

@@ -36,14 +36,14 @@ struct FixedEffectModel <: RegressionModel
     # for FE
     iterations::Int         # Number of iterations
     converged::Bool         # Has the demeaning algorithm converged?
-    r2_within::Union{Float64, Nothing}      # within r2 (with fixed effect
+    r2_within::Float64      # within r2 (with fixed effect
 
     # for IV
-    F_kp::Union{Float64, Nothing}           # First Stage F statistics KP
-    p_kp::Union{Float64, Nothing}           # First Stage p value KP
+    F_kp::Float64           # First Stage F statistics KP
+    p_kp::Float64           # First Stage p value KP
 end
 
-has_iv(m::FixedEffectModel) = m.F_kp !== nothing
+has_iv(m::FixedEffectModel) = has_iv(m.formula)
 has_fe(m::FixedEffectModel) = has_fe(m.formula)
 
 

src/fit.jl

@@ -46,8 +46,9 @@ function partial_out(
         f = FormulaTerm(f.lhs, tuple(ConstantTerm(1), eachterm(f.rhs)...))
     end
     formula, formula_endo, formula_iv = parse_iv(f)
-    has_iv = formula_iv !== nothing
+    has_iv = formula_iv != FormulaTerm(ConstantTerm(0), ConstantTerm(0))
     has_iv && throw("partial_out does not support instrumental variables")
+    formula, formula_fes = parse_fe(formula)
     has_weights = weights !== nothing
 
 
@@ -63,7 +64,7 @@ function partial_out(
     convergeds = Bool[]
 
     # Build fixedeffects, an array of AbtractFixedEffects
-    fes, ids, ids_fes, formula = parse_fixedeffect(df, formula)
+    fes, ids, ids_fes = parse_fixedeffect(df, formula_fes)
     has_fes = !isempty(fes)
 
 

src/partial_out.jl

@@ -4,6 +4,18 @@
 ##
 ##############################################################################
 
+function drop_singletons!(esample, fes::Vector{<:FixedEffect})
+    ns = Int[]
+    for fe in Iterators.cycle(fes)
+        # break loop if number of singletons did not change since the last time fe was iterated on
+        if length(ns) >= length(fes) && sum(view(ns, (length(ns)-length(fes)+1):length(ns))) == ns[end-length(fes)+1]
+            break
+        end
+        push!(ns, drop_singletons!(esample, fe))
+    end
+    return sum(ns)
+end
+
 function drop_singletons!(esample, fe::FixedEffect)
     n = 0
     cache = zeros(Int, fe.n)

src/utils/fixedeffects.jl

@@ -12,23 +12,24 @@ eachterm(@nospecialize(x::NTuple{N, AbstractTerm})) where {N} = x
 ## Parse IV
 ##
 ##############################################################################
-
+has_iv(@nospecialize(f::FormulaTerm)) = any(x -> x isa FormulaTerm, eachterm(f.rhs))
 function parse_iv(@nospecialize(f::FormulaTerm))
-	for term in eachterm(f.rhs)
-		if term isa FormulaTerm
-            both = intersect(eachterm(term.lhs), eachterm(term.rhs))
-            endos = setdiff(eachterm(term.lhs), both)
-            exos = setdiff(eachterm(term.rhs), both)
-            # otherwise empty collection. Later, there will be check
-            isempty(endos) && throw("There are no endogeneous variables")
-            length(exos) < length(endos) && throw("Model not identified. There must be at least as many instrumental variables as endogeneneous variables")
-			formula_endo = FormulaTerm(ConstantTerm(0), tuple(ConstantTerm(0), endos...))
-			formula_iv = FormulaTerm(ConstantTerm(0), tuple(ConstantTerm(0), exos...))
-            formula_exo = FormulaTerm(f.lhs, tuple((term for term in eachterm(f.rhs) if !isa(term, FormulaTerm))..., both...))
-            return formula_exo, formula_endo, formula_iv
-		end
-	end
-	return f, nothing, nothing
+    if has_iv(f)
+        i = findfirst(x -> x isa FormulaTerm, eachterm(f.rhs))
+        term = eachterm(f.rhs)[i]
+        both = intersect(eachterm(term.lhs), eachterm(term.rhs))
+        endos = setdiff(eachterm(term.lhs), both)
+        exos = setdiff(eachterm(term.rhs), both)
+        # otherwise empty collection. Later, there will be check
+        isempty(endos) && throw("There are no endogeneous variables")
+        length(exos) < length(endos) && throw("Model not identified. There must be at least as many instrumental variables as endogeneneous variables")
+		formula_endo = FormulaTerm(ConstantTerm(0), tuple(ConstantTerm(0), endos...))
+		formula_iv = FormulaTerm(ConstantTerm(0), tuple(ConstantTerm(0), exos...))
+        formula_exo = FormulaTerm(f.lhs, tuple((term for term in eachterm(f.rhs) if !isa(term, FormulaTerm))..., both...))
+        return formula_exo, formula_endo, formula_iv
+	else
+    	return f, FormulaTerm(ConstantTerm(0), ConstantTerm(0)), FormulaTerm(ConstantTerm(0), ConstantTerm(0))
+    end
 end
 
 ##############################################################################
@@ -47,13 +48,25 @@ has_fe(::FixedEffectTerm) = true
 has_fe(::FunctionTerm{typeof(fe)}) = true
 has_fe(@nospecialize(t::InteractionTerm)) = any(has_fe(x) for x in t.terms)
 has_fe(::AbstractTerm) = false
-
 has_fe(@nospecialize(t::FormulaTerm)) = any(has_fe(x) for x in eachterm(t.rhs))
 
+function parse_fe(@nospecialize(f::FormulaTerm))
+    if has_fe(f)
+        formula_main = FormulaTerm(f.lhs, Tuple(term for term in eachterm(f.rhs) if !has_fe(term)))
+        formula_fe = FormulaTerm(ConstantTerm(0), Tuple(term for term in eachterm(f.rhs) if has_fe(term)))
+        return formula_main, formula_fe
+    else
+        return f, FormulaTerm(ConstantTerm(0), ConstantTerm(0))
+    end
+end
+
 
 fesymbol(t::FixedEffectTerm) = t.x
 fesymbol(t::FunctionTerm{typeof(fe)}) = Symbol(t.args[1])
 
+
+
+
 """
     parse_fixedeffect(data, formula::FormulaTerm)
     parse_fixedeffect(data, ts::NTuple{N, AbstractTerm})
@@ -68,24 +81,17 @@ Construct any `FixedEffect` specified with a `FixedEffectTerm`.
 """
 function parse_fixedeffect(data, @nospecialize(formula::FormulaTerm))
     fes = FixedEffect[]
-    ids = Symbol[]
+    feids = Symbol[]
     fekeys = Symbol[]
-     for term in eachterm(formula.rhs)
+    for term in eachterm(formula.rhs)
         result = _parse_fixedeffect(data, term)
         if result !== nothing
             push!(fes, result[1])
-            push!(ids, result[2])
+            push!(feids, result[2])
             append!(fekeys, result[3])
         end
     end
-    if !isempty(fes)
-        if any(fe.interaction isa UnitWeights for fe in fes)
-            formula = FormulaTerm(formula.lhs, (InterceptTerm{false}(), (term for term in eachterm(formula.rhs) if !isa(term, Union{ConstantTerm,InterceptTerm}) && !has_fe(term))...))
-        else
-            formula = FormulaTerm(formula.lhs, Tuple(term for term in eachterm(formula.rhs) if !has_fe(term)))
-        end
-    end
-    return fes, ids, unique(fekeys), formula
+    return fes, feids, unique(fekeys)
 end
 
 # Method for external packages
@@ -108,7 +114,7 @@ function parse_fixedeffect(data, @nospecialize(ts::NTuple{N, AbstractTerm})) whe
             ts = Tuple(term for term in eachterm(ts) if !has_fe(term))
         end
     end
-    return fes, ids, unique(fekeys), ts
+    return fes, ids, unique(fekeys)
 end
 
 # Construct FixedEffect from a generic term

src/utils/formula.jl

@@ -22,13 +22,13 @@ for data in [df, csvfile]
     @test _parse_fixedeffect(data, fe(:State)&fe(:Year)) ==
         (FixedEffect(data.State, data.Year), Symbol("fe_State&fe_Year"), [:State, :Year])
 
-    @test parse_fixedeffect(data, ()) == (FixedEffect[], Symbol[], Symbol[], ())
+    @test parse_fixedeffect(data, ()) == (FixedEffect[], Symbol[], Symbol[])
 
     f = @formula(y ~ 1 + Price)
     ts1 = f.rhs
     ts2 = term(1) + term(:Price)
-    @test parse_fixedeffect(data, f) == (FixedEffect[], Symbol[], Symbol[], f)
-    @test parse_fixedeffect(data, ts1) == (FixedEffect[], Symbol[], Symbol[], ts1)
+    @test parse_fixedeffect(data, f) == (FixedEffect[], Symbol[], Symbol[])
+    @test parse_fixedeffect(data, ts1) == (FixedEffect[], Symbol[], Symbol[])
     @test parse_fixedeffect(data, ts2) == parse_fixedeffect(data, ts1)
 
     fparsed = term(:y) ~ InterceptTerm{false}() + term(:Price)
@@ -37,28 +37,28 @@ for data in [df, csvfile]
     f = @formula(y ~ 1 + Price + fe(State))
     ts1 = f.rhs
     ts2 = term(1) + term(:Price) + fe(:State)
-    @test parse_fixedeffect(data, f) == ([FixedEffect(data.State)], [:fe_State], [:State], fparsed)
-    @test parse_fixedeffect(data, ts1) == ([FixedEffect(data.State)], [:fe_State], [:State], tsparsed)
+    @test parse_fixedeffect(data, f) == ([FixedEffect(data.State)], [:fe_State], [:State])
+    @test parse_fixedeffect(data, ts1) == ([FixedEffect(data.State)], [:fe_State], [:State])
     @test parse_fixedeffect(data, ts2) == parse_fixedeffect(data, ts1)
 
     f = @formula(y ~ Price + fe(State) + fe(Year))
     ts1 = f.rhs
     ts2 = term(:Price) + fe(:State) + fe(:Year)
-    @test parse_fixedeffect(data, f) == ([FixedEffect(data.State), FixedEffect(data.Year)], [:fe_State, :fe_Year], [:State, :Year], fparsed)
-    @test parse_fixedeffect(data, ts1) == ([FixedEffect(data.State), FixedEffect(data.Year)], [:fe_State, :fe_Year], [:State, :Year], tsparsed)
+    @test parse_fixedeffect(data, f) == ([FixedEffect(data.State), FixedEffect(data.Year)], [:fe_State, :fe_Year], [:State, :Year])
+    @test parse_fixedeffect(data, ts1) == ([FixedEffect(data.State), FixedEffect(data.Year)], [:fe_State, :fe_Year], [:State, :Year])
     @test parse_fixedeffect(data, ts2) == parse_fixedeffect(data, ts1)
 
     f = @formula(y ~ Price + fe(State)&Year)
     ts1 = f.rhs
     ts2 = term(:Price) + fe(:State)&term(:Year)
-    @test parse_fixedeffect(data, f) == ([FixedEffect(data.State, interaction=_multiply(data, [:Year]))], [Symbol("fe_State&Year")], [:State], term(:y) ~ (term(:Price),))
-    @test parse_fixedeffect(data, ts1) == ([FixedEffect(data.State, interaction=_multiply(data, [:Year]))], [Symbol("fe_State&Year")], [:State], (term(:Price),))
+    @test parse_fixedeffect(data, f) == ([FixedEffect(data.State, interaction=_multiply(data, [:Year]))], [Symbol("fe_State&Year")], [:State])
+    @test parse_fixedeffect(data, ts1) == ([FixedEffect(data.State, interaction=_multiply(data, [:Year]))], [Symbol("fe_State&Year")], [:State])
     @test parse_fixedeffect(data, ts2) == parse_fixedeffect(data, ts1)
 
     f = @formula(y ~ Price + fe(State)*fe(Year))
     ts1 = f.rhs
     ts2 = term(:Price) + fe(:State) + fe(:Year) + fe(:State)&fe(:Year)
-    @test parse_fixedeffect(data, f) == ([FixedEffect(data.State), FixedEffect(data.Year), FixedEffect(data.State, data.Year)], [:fe_State, :fe_Year, Symbol("fe_State&fe_Year")], [:State, :Year], fparsed)
-    @test parse_fixedeffect(data, ts1) == ([FixedEffect(data.State), FixedEffect(data.Year), FixedEffect(data.State, data.Year)], [:fe_State, :fe_Year, Symbol("fe_State&fe_Year")], [:State, :Year], tsparsed)
+    @test parse_fixedeffect(data, f) == ([FixedEffect(data.State), FixedEffect(data.Year), FixedEffect(data.State, data.Year)], [:fe_State, :fe_Year, Symbol("fe_State&fe_Year")], [:State, :Year])
+    @test parse_fixedeffect(data, ts1) == ([FixedEffect(data.State), FixedEffect(data.Year), FixedEffect(data.State, data.Year)], [:fe_State, :fe_Year, Symbol("fe_State&fe_Year")], [:State, :Year])
     @test parse_fixedeffect(data, ts2) == parse_fixedeffect(data, ts1)
 end