Add ParseFEterms, GroupFEterms and MakeFEs (#10)

Author: junyuan-chen

Project.toml

@@ -19,7 +19,7 @@ Vcov = "ec2bfdc2-55df-4fc9-b9ae-4958c2cf2486"
 [compat]
 DataAPI = "1.5"
 DataFrames = "1"
-DiffinDiffsBase = "0.3.2"
+DiffinDiffsBase = "0.3.3"
 FixedEffectModels = "1.5"
 FixedEffects = "2"
 Reexport = "0.2, 1"

src/InteractionWeightedDIDs.jl

@@ -3,7 +3,7 @@ module InteractionWeightedDIDs
 using Base: Callable
 using DataAPI: refarray
 using FixedEffectModels: FixedEffectTerm, Combination,
-    fe, _parse_fixedeffect, invsym!, isnested, nunique
+    fe, fesymbol, _multiply, invsym!, isnested, nunique
 using FixedEffects
 using LinearAlgebra: Cholesky, Factorization, Symmetric, cholesky!, diag
 using Reexport
@@ -26,6 +26,9 @@ export Vcov,
        fe
 
 export CheckVcov,
+       ParseFEterms,
+       GroupFEterms,
+       MakeFEs,
        CheckFEs,
        MakeFESolver,
        MakeYXCols,

src/did.jl

@@ -4,7 +4,8 @@
 Estimation procedure for regression-based difference-in-differences.
 """
 const RegressionBasedDID = DiffinDiffsEstimator{:RegressionBasedDID,
-    Tuple{CheckData, GroupTerms, CheckVcov, CheckVars, CheckFEs, MakeWeights, MakeFESolver,
+    Tuple{CheckData, GroupTreatintterms, GroupXterms, CheckVcov, CheckVars, GroupSample,
+    ParseFEterms, GroupFEterms, MakeFEs, CheckFEs, MakeWeights, MakeFESolver,
     MakeYXCols, MakeTreatCols, SolveLeastSquares, EstVcov, SolveLeastSquaresWeights}}
 
 const Reg = RegressionBasedDID
@@ -64,7 +65,7 @@ struct RegressionBasedDIDResult{TR<:AbstractTreatment, CohortInteracted} <: DIDR
     xterms::Vector{AbstractTerm}
     contrasts::Union{Dict{Symbol, Any}, Nothing}
     weightname::Union{Symbol, Nothing}
-    fenames::Vector{Symbol}
+    fenames::Vector{String}
     nfeiterations::Union{Int, Nothing}
     feconverged::Union{Bool, Nothing}
     nfesingledropped::Int

src/procedures.jl

@@ -26,73 +26,173 @@ exclude rows that are invalid for variance-covariance estimator.
 """
 const CheckVcov = StatsStep{:CheckVcov, typeof(checkvcov!), true}
 
+# Get esample and aux directly from CheckData
 required(::CheckVcov) = (:data, :esample, :aux)
 default(::CheckVcov) = (vce=Vcov.robust(),)
 copyargs(::CheckVcov) = (2,)
 
+"""
+    parsefeterms!(xterms)
+
+Extract any `FixedEffectTerm` or interaction of `FixedEffectTerm` from `xterms`
+and determine whether any intercept term should be omitted.
+See also [`ParseFEterms`](@ref).
+"""
+function parsefeterms!(xterms::TermSet)
+    feterms = Set{FETerm}()
+    has_fe_intercept = false
+    for t in xterms
+        result = _parsefeterm(t)
+        if result !== nothing
+            push!(feterms, result)
+            delete!(xterms, t)
+        end
+    end
+    if !isempty(feterms)
+        if any(t->isempty(t[2]), feterms)
+            has_fe_intercept = true
+            for t in xterms
+                t isa Union{ConstantTerm,InterceptTerm} && delete!(xterms, t)
+            end
+            push!(xterms, InterceptTerm{false}())
+        end
+    end
+    return (xterms=xterms, feterms=feterms, has_fe_intercept=has_fe_intercept)
+end
+
+const ParseFEterms = StatsStep{:ParseFEterms, typeof(parsefeterms!), true}
+
+required(::ParseFEterms) = (:xterms,)
+
+"""
+    groupfeterms(feterms)
+
+Return the argument without change for allowing later comparisons based on object-id.
+See also [`GroupFEterms`](@ref).
+"""
+groupfeterms(feterms::Set{FETerm}) = (feterms=feterms,)
+
+"""
+    GroupFEterms <: StatsStep
+
+Call [`InteractionWeightedDIDs.groupfeterms`](@ref)
+to obtain one of the instances of `feterms`
+that have been grouped by equality (`hash`)
+for allowing later comparisons based on object-id.
+
+This step is only useful when working with [`@specset`](@ref) and [`proceed`](@ref).
+"""
+const GroupFEterms = StatsStep{:GroupFEterms, typeof(groupfeterms), false}
+
+required(::GroupFEterms) = (:feterms,)
+
+"""
+    makefes(args...)
+
+Construct `FixedEffect`s from `data` (the full sample).
+See also [`MakeFEs`](@ref).
+"""
+function makefes(data, allfeterms::Vector{FETerm})
+    # Must use Dict instead of IdDict since the same feterm can be in multiple feterms
+    allfes = Dict{FETerm,FixedEffect}()
+    for t in allfeterms
+        haskey(allfes, t) && continue
+        if isempty(t[2])
+            allfes[t] = FixedEffect((getcolumn(data, n) for n in t[1])...)
+        else
+            allfes[t] = FixedEffect((getcolumn(data, n) for n in t[1])...;
+                interaction=_multiply(data, t[2]))
+        end
+    end
+    return (allfes=allfes,)
+end
+
+"""
+    MakeFEs <: StatsStep
+
+Call [`InteractionWeightedDIDs.makefes`](@ref)
+to construct `FixedEffect`s from `data` (the full sample).
+"""
+const MakeFEs = StatsStep{:MakeFEs, typeof(makefes), false}
+
+required(::MakeFEs) = (:data,)
+combinedargs(::MakeFEs, allntargs) = (FETerm[t for nt in allntargs for t in nt.feterms],)
+
 """
     checkfes!(args...)
 
-Extract any `FixedEffectTerm` from `xterms`,
-drop singleton observations for any fixed effect
-and determine whether intercept term should be omitted.
+Drop any singleton observation from fixed effects over the relevant subsample.
 See also [`CheckFEs`](@ref).
 """
-function checkfes!(data, esample::BitVector, xterms::TermSet, drop_singletons::Bool)
-    fes, fenames, has_fe_intercept = parse_fixedeffect!(data, xterms)
+function checkfes!(feterms::Set{FETerm}, allfes::Dict{FETerm,FixedEffect},
+        esample::BitVector, drop_singletons::Bool)
     nsingle = 0
-    if !isempty(fes)
+    nfe = length(feterms)
+    if nfe > 0
+        fes = Vector{FixedEffect}(undef, nfe)
+        fenames = Vector{String}(undef, nfe)
+        # Loop together to ensure the orders are the same
+        for (i, t) in enumerate(feterms)
+            fes[i] = allfes[t]
+            fenames[i] = getfename(t)
+        end
+        # Determine the unique order based on names
+        order = sortperm(fenames)
+        fes = fes[order]
+        fenames = fenames[order]
+
         if drop_singletons
             for fe in fes
                 nsingle += drop_singletons!(esample, fe)
             end
         end
+        sum(esample) == 0 && error("no nonmissing data")
+
+        for i in 1:nfe
+            fes[i] = fes[i][esample]
+        end
+        return (esample=esample, fes=fes, fenames=fenames, nsingle=nsingle)
+    else
+        return (esample=esample, fes=FixedEffect[], fenames=String[], nsingle=0)
     end
-    sum(esample) == 0 && error("no nonmissing data")
-    return (xterms=xterms, esample=esample, fes=fes, fenames=fenames,
-        has_fe_intercept=has_fe_intercept, nsingle=nsingle)
 end
 
 """
     CheckFEs <: StatsStep
 
 Call [`InteractionWeightedDIDs.checkfes!`](@ref)
-to extract any `FixedEffectTerm` from `xterms`
-and drop singleton observations for any fixed effect.
+to drop any singleton observation from fixed effects over the relevant subsample.
 """
-const CheckFEs = StatsStep{:CheckFixedEffects, typeof(checkfes!), true}
+const CheckFEs = StatsStep{:CheckFEs, typeof(checkfes!), true}
 
-required(::CheckFEs) = (:data, :esample, :xterms)
+required(::CheckFEs) = (:feterms, :allfes, :esample)
 default(::CheckFEs) = (drop_singletons=true,)
-copyargs(::CheckFEs) = (2,3)
+copyargs(::CheckFEs) = (3,)
 
 """
-    makefesolver!(args...)
+    makefesolver(args...)
 
 Construct `FixedEffects.AbstractFixedEffectSolver`.
 See also [`MakeFESolver`](@ref).
 """
-function makefesolver!(fes::Vector{FixedEffect}, weights::AbstractWeights,
-        esample::BitVector, nfethreads::Int)
+function makefesolver(fes::Vector{FixedEffect}, weights::AbstractWeights, nfethreads::Int)
     if !isempty(fes)
-        fes = FixedEffect[fe[esample] for fe in fes]
         feM = AbstractFixedEffectSolver{Float64}(fes, weights, Val{:cpu}, nfethreads)
-        return (feM=feM, fes=fes)
+        return (feM=feM,)
     else
-        return (feM=nothing, fes=fes)
+        return (feM=nothing,)
     end
 end
 
 """
     MakeFESolver <: StatsStep
 
-Call [`InteractionWeightedDIDs.makefesolver!`](@ref) to construct the fixed effect solver.
+Call [`InteractionWeightedDIDs.makefesolver`](@ref) to construct the fixed effect solver.
 """
-const MakeFESolver = StatsStep{:MakeFESolver, typeof(makefesolver!), true}
+const MakeFESolver = StatsStep{:MakeFESolver, typeof(makefesolver), true}
 
-required(::MakeFESolver) = (:fes, :weights, :esample)
+required(::MakeFESolver) = (:fes, :weights)
 default(::MakeFESolver) = (nfethreads=Threads.nthreads(),)
-copyargs(::MakeFESolver) = (1,)
 
 function _makeyxcols!(yxterms::Dict, yxcols::Dict, yxschema, data, t::AbstractTerm)
     ct = apply_schema(t, yxschema, StatisticalModel)
@@ -291,6 +391,7 @@ transformed(::MakeTreatCols, @nospecialize(nt::NamedTuple)) = (nt.tr.time,)
 # Obtain the relative time periods excluded by all tr
 # and the treatment groups excluded by all pr in allntargs
 function combinedargs(::MakeTreatCols, allntargs)
+    # exc cannot be IdDict for comparing different types of one
     exc = Dict{Int,Int}()
     notreat = IdDict{ValidTimeType,Int}()
     for nt in allntargs

src/utils.jl

@@ -1,30 +1,27 @@
-function parse_fixedeffect!(data, ts::TermSet)
-    fes = FixedEffect[]
-    ids = Symbol[]
-    has_fe_intercept = false
-    for term in ts
-        result = _parse_fixedeffect(data, term)
-        if result !== nothing
-            push!(fes, result[1])
-            push!(ids, result[2])
-            delete!(ts, term)
-        end
+# A vector of fixed effects paired with a vector of interactions (empty if not interacted)
+const FETerm = Pair{Vector{Symbol},Vector{Symbol}}
+
+# Parse fixed effects from a generic term
+function _parsefeterm(@nospecialize(t::AbstractTerm))
+    if has_fe(t)
+        s = fesymbol(t)
+        return [s]=>Symbol[]
     end
-    order = sortperm(ids)
-    fes .= fes[order]
-    ids .= ids[order]
+end
+
+# Parse fixed effects from an InteractionTerm
+function _parsefeterm(@nospecialize(t::InteractionTerm))
+    fes = (x for x in t.terms if has_fe(x))
+    interactions = (x for x in t.terms if !has_fe(x))
     if !isempty(fes)
-        if any(fe->fe.interaction isa UnitWeights, fes)
-            has_fe_intercept = true
-            for t in ts
-                t isa Union{ConstantTerm,InterceptTerm} && delete!(ts, t)
-            end
-            push!(ts, InterceptTerm{false}())
-        end
+        fes = sort!([fesymbol(x) for x in fes])
+        feints = sort!([Symbol(x) for x in interactions])
+        return fes=>feints
     end
-    return fes, ids, has_fe_intercept
 end
 
+getfename(feterm::FETerm) = join(vcat("fe_".*string.(feterm[1]), feterm[2]), "&")
+
 # Count the number of singletons dropped
 function drop_singletons!(esample, fe::FixedEffect)
     cache = zeros(Int, fe.n)

test/did.jl

@@ -209,19 +209,25 @@ end
 @testset "@specset" begin
     hrs = exampledata("hrs")
     # The first two specs are identical hence no repetition of steps should occur
-    # The third spec should only share the first three steps with the others
-    r = @specset [verbose] begin
-        @did(Reg, dynamic(:wave, -1), notyettreated(11), data=hrs,
-            yterm=term(:oop_spend), treatname=:wave_hosp, treatintterms=(),
+    # The third spec should share all the steps until SolveLeastSquares
+    # The fourth and fifth specs should not add tasks for MakeFEs and MakeYXCols
+    # The sixth spec should not add any task for MakeFEs
+    r = @specset [verbose] data=hrs yterm=term(:oop_spend) treatname=:wave_hosp begin
+        @did(Reg, dynamic(:wave, -1), notyettreated(11),
             xterms=(fe(:wave)+fe(:hhidpn)))
-        @did(Reg, dynamic(:wave, -1), notyettreated(11), data=hrs,
-            yterm=term(:oop_spend), treatname=:wave_hosp, treatintterms=[],
+        @did(Reg, dynamic(:wave, -1), notyettreated(11),
             xterms=[fe(:hhidpn), fe(:wave)])
-        @did(Reg, dynamic(:wave, -1), nevertreated(11), data=hrs,
-            yterm=term(:oop_spend), treatname=:wave_hosp, treatintterms=(),
+        @did(Reg, dynamic(:wave, -2:-1), notyettreated(11),
+            xterms=[fe(:hhidpn), fe(:wave)])
+        @did(Reg, dynamic(:wave, -1), notyettreated(11),
+            xterms=[term(:male), fe(:hhidpn), fe(:wave)])
+        @did(Reg, dynamic(:wave, -1), notyettreated(11),
+            treatintterms=TermSet(:male), xterms=[fe(:hhidpn), fe(:wave)])
+        @did(Reg, dynamic(:wave, -1), nevertreated(11),
             xterms=(fe(:wave)+fe(:hhidpn)))
     end
-    @test r[1] == didspec(Reg, dynamic(:wave, -1), notyettreated(11), data=hrs,
+    # Results might differ due to yxterms that include terms from other specs
+    @test r[4] == didspec(Reg, dynamic(:wave, -1), notyettreated(11), data=hrs,
         yterm=term(:oop_spend), treatname=:wave_hosp, treatintterms=(),
-        xterms=TermSet(fe(:wave), fe(:hhidpn)))()
+        xterms=TermSet(term(:male), fe(:wave), fe(:hhidpn)))()
 end