Add SolveLeastSquaresWeights, AggregatedRegBasedDIDResult and agg (#6)

Author: junyuan-chen

Project.toml

@@ -16,7 +16,8 @@ Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 Vcov = "ec2bfdc2-55df-4fc9-b9ae-4958c2cf2486"
 
 [compat]
-DiffinDiffsBase = "0.2.1"
+DataFrames = "0.22"
+DiffinDiffsBase = "0.3"
 FixedEffectModels = "1.4"
 FixedEffects = "2"
 Reexport = "0.2, 1"

src/InteractionWeightedDIDs.jl

@@ -13,11 +13,12 @@ using Tables
 using Tables: getcolumn, columnnames
 using Vcov
 @reexport using DiffinDiffsBase
-using DiffinDiffsBase: TimeType, termvars, isintercept, parse_intercept!, _treatnames
+using DiffinDiffsBase: TimeType, termvars, isintercept, parse_intercept!,
+    _count!, _groupfind, _treatnames, _parse_bycells!, _parse_subset
 
 import Base: show
 import DiffinDiffsBase: required, default, transformed, combinedargs, copyargs,
-    valid_didargs, result, _count!
+    valid_didargs, result, vce, nobs, outcomename, weights, treatnames, dof_residual, agg
 import FixedEffectModels: has_fe
 
 export Vcov,
@@ -30,10 +31,12 @@ export CheckVcov,
        MakeTreatCols,
        SolveLeastSquares,
        EstVcov,
+       SolveLeastSquaresWeights,
 
        RegressionBasedDID,
        Reg,
-       RegressionBasedDIDResult
+       RegressionBasedDIDResult,
+       AggregatedRegBasedDIDResult
 
 include("utils.jl")
 include("procedures.jl")

src/did.jl

@@ -5,7 +5,7 @@ Estimation procedure for regression-based difference-in-differences.
 """
 const RegressionBasedDID = DiffinDiffsEstimator{:RegressionBasedDID,
     Tuple{CheckData, GroupTerms, CheckVcov, CheckVars, CheckFEs, MakeWeights, MakeFESolver,
-    MakeYXCols, MakeTreatCols, SolveLeastSquares, EstVcov}}
+    MakeYXCols, MakeTreatCols, SolveLeastSquares, EstVcov, SolveLeastSquaresWeights}}
 
 const Reg = RegressionBasedDID
 
@@ -14,6 +14,7 @@ function valid_didargs(d::Type{Reg}, ::DynamicTreatment{SharpDesign},
     name = get(args, :name, "")::String
     treatintterms = haskey(args, :treatintterms) ? args[:treatintterms] : TermSet()
     xterms = haskey(args, :xterms) ? args[:xterms] : TermSet()
+    solvelsweights = haskey(args, :lswtnames) || get(args, :solvelsweights, false)::Bool
     ntargs = (data=args[:data],
         tr=args[:tr]::DynamicTreatment{SharpDesign},
         pr=args[:pr]::TrendParallel{Unconditional, Exact},
@@ -29,15 +30,22 @@ function valid_didargs(d::Type{Reg}, ::DynamicTreatment{SharpDesign},
         contrasts=get(args, :contrasts, nothing)::Union{Dict{Symbol,Any},Nothing},
         fetol=get(args, :fetol, 1e-8)::Float64,
         femaxiter=get(args, :femaxiter, 10000)::Int,
-        cohortinteracted=get(args, :cohortinteracted, true)::Bool)
+        cohortinteracted=get(args, :cohortinteracted, true)::Bool,
+        solvelsweights=solvelsweights::Bool,
+        lswtnames=get(args, :lswtnames, ()))
     return name, d, ntargs
 end
 
-struct RegressionBasedDIDResult{CohortInteracted} <: DIDResult
+"""
+    RegressionBasedDIDResult{TR<:AbstractTreatment, CohortInteracted} <: DIDResult
+
+Estimation results from regression-based difference-in-differences.
+"""
+struct RegressionBasedDIDResult{TR<:AbstractTreatment, CohortInteracted} <: DIDResult
     coef::Vector{Float64}
     vcov::Matrix{Float64}
     vce::CovarianceEstimator
-    tr::AbstractTreatment
+    tr::TR
     pr::AbstractParallel
     cellweights::Vector{Float64}
     cellcounts::Vector{Int}
@@ -60,6 +68,10 @@ struct RegressionBasedDIDResult{CohortInteracted} <: DIDResult
     nfeiterations::Union{Int, Nothing}
     feconverged::Union{Bool, Nothing}
     nfesingledropped::Int
+    lsweights::Union{TableIndexedMatrix{Float64, Matrix{Float64}, VecColumnTable, VecColumnTable}, Nothing}
+    ycellmeans::Union{Vector{Float64}, Nothing}
+    ycellweights::Union{Vector{Float64}, Nothing}
+    ycellcounts::Union{Vector{Int}, Nothing}
 end
 
 function result(::Type{Reg}, @nospecialize(nt::NamedTuple))
@@ -68,12 +80,13 @@ function result(::Type{Reg}, @nospecialize(nt::NamedTuple))
     cnames = _treatnames(nt.treatcells)
     cnames = append!(cnames, coefnames.(nt.xterms))[nt.basecols]
     coefinds = Dict(cnames .=> 1:length(cnames))
-    didresult = RegressionBasedDIDResult{nt.cohortinteracted}(
+    didresult = RegressionBasedDIDResult{typeof(nt.tr), nt.cohortinteracted}(
         nt.coef, nt.vcov_mat, nt.vce, nt.tr, nt.pr, nt.cellweights, nt.cellcounts,
         nt.esample, sum(nt.esample), nt.dof_resid, nt.F, nt.p,
         yname, cnames, coefinds, nt.treatcells, nt.treatname, nt.yxterms,
         yterm, nt.xterms, nt.contrasts, nt.weightname,
-        nt.fenames, nt.nfeiterations, nt.feconverged, nt.nsingle)
+        nt.fenames, nt.nfeiterations, nt.feconverged, nt.nsingle,
+        nt.lsweights, nt.ycellmeans, nt.ycellweights, nt.ycellcounts)
     return merge(nt, (result=didresult,))
 end
 
@@ -96,22 +109,22 @@ _nunique(t, s::Symbol) = length(unique(getproperty(t, s)))
 _excluded_rel_str(tr::DynamicTreatment) =
     isempty(tr.exc) ? "none" : join(string.(tr.exc), " ")
 
-_treat_info(r::RegressionBasedDIDResult{true}, tr::DynamicTreatment) = (
+_treat_info(r::RegressionBasedDIDResult{<:DynamicTreatment, true}) = (
     "Number of cohorts" => _nunique(r.treatcells, r.treatname),
     "Interactions within cohorts" => length(columnnames(r.treatcells)) - 2,
     "Relative time periods" => _nunique(r.treatcells, :rel),
-    "Excluded periods" => NoQuote(_excluded_rel_str(tr))
+    "Excluded periods" => NoQuote(_excluded_rel_str(r.tr))
 )
 
-_treat_info(r::RegressionBasedDIDResult{false}, tr::DynamicTreatment) = (
+_treat_info(r::RegressionBasedDIDResult{<:DynamicTreatment, false}) = (
     "Relative time periods" => _nunique(r.treatcells, :rel),
-    "Excluded periods" => NoQuote(_excluded_rel_str(tr))
+    "Excluded periods" => NoQuote(_excluded_rel_str(r.tr))
 )
 
-_treat_spec(r::RegressionBasedDIDResult{true}, tr::DynamicTreatment{SharpDesign}) =
+_treat_spec(r::RegressionBasedDIDResult{DynamicTreatment{SharpDesign}, true}) =
     "Cohort-interacted sharp dynamic specification"
 
-_treat_spec(r::RegressionBasedDIDResult{false}, tr::DynamicTreatment{SharpDesign}) =
+_treat_spec(r::RegressionBasedDIDResult{DynamicTreatment{SharpDesign}, false}) =
     "Sharp dynamic specification"
 
 _fe_info(r::RegressionBasedDIDResult) = (
@@ -126,12 +139,12 @@ function show(io::IO, ::MIME"text/plain", r::RegressionBasedDIDResult;
     halfwidth = div(totalwidth-interwidth, 2)
     top_info = _top_info(r)
     fe_info = has_fe(r) ? _fe_info(r) : ()
-    tr_info = _treat_info(r, r.tr)
+    tr_info = _treat_info(r)
     blocks = (top_info, tr_info, fe_info)
     fes = has_fe(r) ? join(string.(r.fenames), " ") : "none"
     fetitle = string("Fixed effects: ", fes)
     blocktitles = ("Summary of results: Regression-based DID",
-        _treat_spec(r, r.tr), fetitle[1:min(totalwidth,length(fetitle))])
+        _treat_spec(r), fetitle[1:min(totalwidth,length(fetitle))])
 
     for (ib, b) in enumerate(blocks)
         println(io, repeat('─', totalwidth))
@@ -148,3 +161,77 @@ function show(io::IO, ::MIME"text/plain", r::RegressionBasedDIDResult;
     end
     print(io, repeat('─', totalwidth))
 end
+
+"""
+    AggregatedRegBasedDIDResult{P<:RegressionBasedDIDResult, I} <: AggregatedDIDResult{P}
+
+Estimation results aggregated from a [`RegressionBasedDIDResult`](@ref).
+See also [`agg`](@ref).
+"""
+struct AggregatedRegBasedDIDResult{P<:RegressionBasedDIDResult, I} <: AggregatedDIDResult{P}
+    parent::P
+    inds::I
+    coef::Vector{Float64}
+    vcov::Matrix{Float64}
+    coefweights::Matrix{Float64}
+    cellweights::Vector{Float64}
+    cellcounts::Vector{Int}
+    coefnames::Vector{String}
+    coefinds::Dict{String, Int}
+    treatcells::VecColumnTable
+    lsweights::Union{TableIndexedMatrix{Float64, Matrix{Float64}, VecColumnTable, VecColumnTable}, Nothing}
+end
+
+function agg(r::RegressionBasedDIDResult{<:DynamicTreatment}, names=nothing;
+        bys=nothing, subset=nothing)
+    inds = subset === nothing ? Colon() : _parse_subset(r, subset, false)
+    ptcells = treatcells(r)
+    bycells = view(ptcells, inds)
+    _parse_bycells!(getfield(bycells, :columns), ptcells, bys)
+    names === nothing || (bycells = subcolumns(bycells, names, nomissing=false))
+
+    tcells, rows = cellrows(bycells, findcell(bycells))
+    ncell = length(rows)
+    pcf = view(treatcoef(r), inds)
+    cweights = zeros(length(pcf), ncell)
+    pcellweights = view(r.cellweights, inds)
+    pcellcounts = view(r.cellcounts, inds)
+    # Ensure the weights for each relative time always sum up to one
+    rels = view(r.treatcells.rel, inds)
+    for (i, rs) in enumerate(rows)
+        if length(rs) > 1
+            relgroups = _groupfind(view(rels, rs))
+            for inds in values(relgroups)
+                if length(inds) > 1
+                    cwts = view(pcellweights, view(rs, inds))
+                    cweights[view(rs, inds), i] .= cwts ./ sum(cwts)
+                else
+                    cweights[rs[inds[1]], i] = 1.0
+                end
+            end
+        else
+            cweights[rs[1], i] = 1.0
+        end
+    end
+    cf = cweights' * pcf
+    v = cweights' * view(treatvcov(r), inds, inds) * cweights
+    cellweights = [sum(pcellweights[rows[i]]) for i in 1:ncell]
+    cellcounts = [sum(pcellcounts[rows[i]]) for i in 1:ncell]
+    cnames = _treatnames(tcells)
+    coefinds = Dict(cnames .=> keys(cnames))
+    if r.lsweights === nothing
+        lswt = nothing
+    else
+        lswtmat = view(r.lsweights.m, :, inds) * cweights
+        lswt = TableIndexedMatrix(lswtmat, r.lsweights.r, tcells)
+    end
+    return AggregatedRegBasedDIDResult{typeof(r), typeof(inds)}(r, inds, cf, v, cweights,
+        cellweights, cellcounts, cnames, coefinds, tcells, lswt)
+end
+
+vce(r::AggregatedRegBasedDIDResult) = vce(parent(r))
+nobs(r::AggregatedRegBasedDIDResult) = nobs(parent(r))
+outcomename(r::AggregatedRegBasedDIDResult) = outcomename(parent(r))
+weights(r::AggregatedRegBasedDIDResult) = weights(parent(r))
+treatnames(r::AggregatedRegBasedDIDResult) = coefnames(r)
+dof_residual(r::AggregatedRegBasedDIDResult) = dof_residual(parent(r))