Add GroupTreatintterms, GroupXterms, GroupSample and post! (#32)

Author: junyuan-chen

Project.toml

@@ -35,7 +35,7 @@ Reexport = "0.2, 1"
 StatsBase = "0.33"
 StatsFuns = "0.9"
 StatsModels = "0.6.18"
-StructArrays = "0.5"
+StructArrays = "0.5, 0.6"
 Tables = "1.2"
 julia = "1.3"
 

src/DiffinDiffsBase.jl

@@ -109,8 +109,10 @@ export cb,
        @specset,
 
        CheckData,
-       GroupTerms,
+       GroupTreatintterms,
+       GroupXterms,
        CheckVars,
+       GroupSample,
        MakeWeights,
 
        DiffinDiffsEstimator,
@@ -135,7 +137,12 @@ export cb,
        TransformedDIDResult,
        TransSubDIDResult,
        lincom,
-       rescale
+       rescale,
+       ExportFormat,
+       StataPostHDF,
+       getexportformat,
+       setexportformat!,
+       post!
 
 include("tables.jl")
 include("utils.jl")

src/did.jl

@@ -780,3 +780,67 @@ function rescale(r::AbstractDIDResult, by::Pair, subset)
     tinds = treatindex(ntreatcoef(r), inds)
     return rescale(r, apply(view(treatcells(r), tinds), by), inds)
 end
+
+"""
+    ExportFormat
+
+Supertype for all types representing the format for exporting an [`AbstractDIDResult`](@ref).
+"""
+abstract type ExportFormat end
+
+"""
+    StataPostHDF <: ExportFormat
+
+Export an [`AbstractDIDResult`](@ref) for Stata module
+[`posthdf`](https://github.com/junyuan-chen/posthdf).
+"""
+struct StataPostHDF <: ExportFormat end
+
+const DefaultExportFormat = ExportFormat[StataPostHDF()]
+
+"""
+    getexportformat()
+
+Return the default [`ExportFormat`](@ref) for [`post!`](@ref).
+"""
+getexportformat() = DefaultExportFormat[1]
+
+"""
+    setexportformat!(format::ExportFormat)
+
+Set the default [`ExportFormat`](@ref) for [`post!`](@ref).
+"""
+setexportformat!(format::ExportFormat) = (DefaultExportFormat[1] = format)
+
+"""
+    post!(f, r::AbstractDIDResult; kwargs...)
+
+Export result `r` in a default [`ExportFormat`](@ref).
+
+The default format can be retrieved via [`getexportformat`](@ref)
+and modified via [`setexportformat!`](@ref).
+"""
+post!(f, r::AbstractDIDResult; kwargs...) =
+    post!(f, getexportformat(), r; kwargs...)
+
+"""
+    post!(f, ::StataPostHDF, r::AbstractDIDResult; model="DiffinDiffsBase.AbstractDIDResult")
+
+Export result `r` for Stata module
+[`posthdf`](https://github.com/junyuan-chen/posthdf).
+A subset of field values from `r` are placed in `f` by setting key-value pairs,
+where `f` can be either an `HDF5.Group` or any object that can be indexed by strings.
+"""
+function post!(f, ::StataPostHDF, r::AbstractDIDResult;
+        model::String="DiffinDiffsBase.AbstractDIDResult")
+    f["model"] = model
+    f["b"] = coef(r)
+    f["V"] = vcov(r)
+    f["vce"] = repr(vce(r))
+    f["N"] = nobs(r)
+    f["depvar"] = string(outcomename(r))
+    f["coefnames"] = convert(AbstractVector{String}, coefnames(r))
+    f["weights"] = (w = weights(r); w === nothing ? "" : string(w))
+    f["ntreatcoef"] = ntreatcoef(r)
+    return f
+end

src/procedures.jl

@@ -45,27 +45,48 @@ required(::CheckData) = (:data,)
 default(::CheckData) = (subset=nothing, weightname=nothing)
 
 """
-    groupterms(args...)
+    grouptreatintterms(treatintterms)
 
-Return the arguments for allowing later comparisons based on object-id.
-See also [`GroupTerms`](@ref).
+Return the argument without change for allowing later comparisons based on object-id.
+See also [`GroupTreatintterms`](@ref).
 """
-groupterms(treatintterms::TermSet, xterms::TermSet) =
-    (treatintterms = treatintterms, xterms = xterms)
+grouptreatintterms(treatintterms::TermSet) = (treatintterms=treatintterms,)
 
 """
-    GroupTerms <: StatsStep
+    GroupTreatintterms <: StatsStep
 
-Call [`DiffinDiffsBase.groupterms`](@ref)
-to obtain one of the instances of `treatintterms` and `xterms`
-that have been grouped by `==`
+Call [`DiffinDiffsBase.grouptreatintterms`](@ref)
+to obtain one of the instances of `treatintterms`
+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 GroupTerms = StatsStep{:GroupTerms, typeof(groupterms), false}
+const GroupTreatintterms = StatsStep{:GroupTreatintterms, typeof(grouptreatintterms), false}
 
-required(::GroupTerms) = (:treatintterms, :xterms)
+default(::GroupTreatintterms) = (treatintterms=TermSet(),)
+
+"""
+    groupxterms(xterms)
+
+Return the argument without change for allowing later comparisons based on object-id.
+See also [`GroupXterms`](@ref).
+"""
+groupxterms(xterms::TermSet) = (xterms=xterms,)
+
+"""
+    GroupXterms <: StatsStep
+
+Call [`DiffinDiffsBase.groupxterms`](@ref)
+to obtain one of the instances of `xterms`
+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 GroupXterms = StatsStep{:GroupXterms, typeof(groupxterms), false}
+
+default(::GroupXterms) = (xterms=TermSet(),)
 
 function _checkscales(col1::AbstractArray, col2::AbstractArray, treatvars::Vector{Symbol})
     if col1 isa ScaledArrOrSub || col2 isa ScaledArrOrSub
@@ -159,11 +180,12 @@ Exclude rows with missing data or violate the overlap condition
 and find rows with data from treated units.
 See also [`CheckVars`](@ref).
 """
-function checkvars!(data, tr::AbstractTreatment, pr::AbstractParallel,
+function checkvars!(data, pr::AbstractParallel,
         yterm::AbstractTerm, treatname::Symbol, esample::BitVector, aux::BitVector,
-        treatintterms::TermSet, xterms::TermSet)
+        treatintterms::TermSet, xterms::TermSet, ::Type, @nospecialize(trvars::Tuple),
+        tr::AbstractTreatment)
     # Do not check eltype of treatintterms
-    treatvars = union([treatname], termvars(tr), termvars(pr))
+    treatvars = union([treatname], trvars, termvars(pr))
     checktreatvars(tr, pr, treatvars, data)
 
     allvars = union(treatvars, termvars(yterm), termvars(xterms))
@@ -202,9 +224,40 @@ Call [`DiffinDiffsBase.checkvars!`](@ref) to exclude invalid rows for relevant v
 """
 const CheckVars = StatsStep{:CheckVars, typeof(checkvars!), true}
 
-required(::CheckVars) = (:data, :tr, :pr, :yterm, :treatname, :esample, :aux)
-default(::CheckVars) = (treatintterms=TermSet(), xterms=TermSet())
-copyargs(::CheckVars) = (6,)
+required(::CheckVars) = (:data, :pr, :yterm, :treatname, :esample, :aux,
+    :treatintterms, :xterms)
+transformed(::CheckVars, @nospecialize(nt::NamedTuple)) =
+    (typeof(nt.tr), (termvars(nt.tr)...,))
+
+combinedargs(step::CheckVars, allntargs) =
+    combinedargs(step, allntargs, typeof(allntargs[1].tr))
+
+combinedargs(::CheckVars, allntargs, ::Type{DynamicTreatment{SharpDesign}}) =
+    (allntargs[1].tr,)
+
+copyargs(::CheckVars) = (5,)
+
+"""
+    groupsample(esample)
+
+Return the argument without change for allowing later comparisons based on object-id.
+See also [`GroupSample`](@ref).
+"""
+groupsample(esample::BitVector) = (esample=esample,)
+
+"""
+    GroupSample <: StatsStep
+
+Call [`DiffinDiffsBase.groupsample`](@ref)
+to obtain one of the instances of `esample`
+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 GroupSample = StatsStep{:GroupSample, typeof(groupsample), false}
+
+required(::GroupSample) = (:esample,)
 
 """
     makeweights(args...)

test/did.jl

@@ -545,3 +545,22 @@ end
     m = Diagonal(r.treatcells.rel[1:3])
     @test vcov(tr) == m * r.vcov[1:3,1:3] * m'
 end
+
+@testset "post!" begin
+    @test getexportformat() == DefaultExportFormat[1]
+    setexportformat!(StataPostHDF())
+    @test DefaultExportFormat[1] == StataPostHDF()
+
+    f = Dict{String,Any}()
+    r = TestResult(2, 2)
+    post!(f, r)
+    @test f["model"] == "DiffinDiffsBase.AbstractDIDResult"
+    @test f["b"] == coef(r)
+    @test f["V"] == vcov(r)
+    @test f["vce"] == "nothing"
+    @test f["N"] == nobs(r)
+    @test f["depvar"] == "y"
+    @test f["coefnames"][1] == "rel: 1 & c: 1"
+    @test f["weights"] == "w"
+    @test f["ntreatcoef"] == 4
+end

test/procedures.jl

@@ -42,19 +42,35 @@
     end
 end
 
-@testset "GroupTerms" begin
-    @testset "groupterms" begin
-        nt = (treatintterms=TermSet(), xterms=TermSet(term(:x)))
-        @test groupterms(nt...) == nt
+@testset "GroupTreatintterms" begin
+    @testset "grouptreatintterms" begin
+        nt = (treatintterms=TermSet(),)
+        @test grouptreatintterms(nt...) == nt
     end
 
     @testset "StatsStep" begin
-        @test sprint(show, GroupTerms()) == "GroupTerms"
-        @test sprint(show, MIME("text/plain"), GroupTerms()) ==
-            "GroupTerms (StatsStep that calls DiffinDiffsBase.groupterms)"
-        @test _byid(GroupTerms()) == false
-        nt = (treatintterms=TermSet(), xterms=TermSet(term(:x)))
-        @test GroupTerms()(nt) == nt 
+        @test sprint(show, GroupTreatintterms()) == "GroupTreatintterms"
+        @test sprint(show, MIME("text/plain"), GroupTreatintterms()) ==
+            "GroupTreatintterms (StatsStep that calls DiffinDiffsBase.grouptreatintterms)"
+        @test _byid(GroupTreatintterms()) == false
+        nt = (treatintterms=TermSet(),)
+        @test GroupTreatintterms()() == nt
+    end
+end
+
+@testset "GroupXterms" begin
+    @testset "groupxterms" begin
+        nt = (xterms=TermSet(term(:x)),)
+        @test groupxterms(nt...) == nt
+    end
+
+    @testset "StatsStep" begin
+        @test sprint(show, GroupXterms()) == "GroupXterms"
+        @test sprint(show, MIME("text/plain"), GroupXterms()) ==
+            "GroupXterms (StatsStep that calls DiffinDiffsBase.groupxterms)"
+        @test _byid(GroupXterms()) == false
+        nt = (xterms=TermSet(),)
+        @test GroupXterms()() == nt
     end
 end
 
@@ -63,9 +79,11 @@ end
         hrs = exampledata("hrs")
         N = size(hrs,1)
         us = unspecifiedpr()
-        nt = (data=hrs, tr=dynamic(:wave, -1), pr=us, yterm=term(:oop_spend),
+        tr = dynamic(:wave, -1)
+        nt = (data=hrs, pr=us, yterm=term(:oop_spend),
             treatname=:wave_hosp, esample=trues(N), aux=BitVector(undef, N),
-            treatintterms=TermSet(), xterms=TermSet())
+            treatintterms=TermSet(), xterms=TermSet(),
+            tytr=typeof(tr), trvars=(termvars(tr)...,), tr=tr)
         @test checkvars!(nt...) == (esample=trues(N), tr_rows=trues(N))
 
         nt = merge(nt, (pr=nevertreated(11),))
@@ -140,7 +158,8 @@ end
         df.wave_hosp = rotatingtime(rot, df.wave_hosp)
         df.wave = rotatingtime(rot, df.wave)
         e = rotatingtime((1,2), 11)
-        nt = merge(nt, (data=df, tr=dynamic(:wave, -1), pr=nevertreated(e), treatintterms=TermSet(), xterms=TermSet(), esample=trues(N)))
+        nt = merge(nt, (data=df, pr=nevertreated(e), treatintterms=TermSet(),
+            xterms=TermSet(), esample=trues(N)))
         # Check RotatingTimeArray
         @test_throws ArgumentError checkvars!(nt...)
         df.wave_hosp = settime(hrs.wave_hosp, rotation=rot)
@@ -198,13 +217,30 @@ end
         @test CheckVars()(nt) ==
             merge(nt, (esample=trues(N), tr_rows=hrs.wave_hosp.!=11))
         nt = (data=hrs, tr=dynamic(:wave, -1), pr=nevertreated(11), yterm=term(:oop_spend),
-            treatname=:wave_hosp, esample=trues(N), aux=BitVector(undef, N))
+            treatname=:wave_hosp, esample=trues(N), aux=BitVector(undef, N),
+            treatintterms=TermSet(), xterms=TermSet())
         @test CheckVars()(nt) ==
             merge(nt, (esample=trues(N), tr_rows=hrs.wave_hosp.!=11))
         @test_throws ErrorException CheckVars()()
     end
 end
 
+@testset "GroupSample" begin
+    @testset "groupsample" begin
+        nt = (esample=trues(3),)
+        @test groupsample(nt...) == nt
+    end
+
+    @testset "StatsStep" begin
+        @test sprint(show, GroupSample()) == "GroupSample"
+        @test sprint(show, MIME("text/plain"), GroupSample()) ==
+            "GroupSample (StatsStep that calls DiffinDiffsBase.groupsample)"
+        @test _byid(GroupSample()) == false
+        nt = (esample=trues(3),)
+        @test GroupSample()(nt) == nt
+    end
+end
+
 @testset "MakeWeights" begin
     @testset "makeweights" begin
         hrs = exampledata("hrs")

test/runtests.jl

@@ -7,9 +7,10 @@ using Dates: Date, Year
 using DiffinDiffsBase: @fieldequal, unpack, @unpack, checktable, hastreat, parse_treat,
     isintercept, isomitsintercept, parse_intercept!,
     ncol, nrow, _mult!,
-    _f, _byid, groupargs, copyargs, pool, checkdata!, groupterms, checkvars!, makeweights,
+    _f, _byid, groupargs, copyargs, pool,
+    checkdata!, grouptreatintterms, groupxterms, checkvars!, groupsample, makeweights,
     _totermset!, parse_didargs!, _treatnames, _parse_bycells!, _parse_subset, _nselected,
-    treatindex, checktreatindex
+    treatindex, checktreatindex, DefaultExportFormat
 using LinearAlgebra: Diagonal
 using Missings: allowmissing, disallowmissing
 using PooledArrays: PooledArray