Improve StatsStep and add CheckVars

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -48,6 +48,7 @@ export @fieldequal,
        nevertreated,
        NotYetTreatedParallel,
        notyettreated,
+       treated,
 
        TreatmentTerm,
        treat,
@@ -65,6 +66,7 @@ export @fieldequal,
        proceed,
 
        CheckData,
+       CheckVars,
 
        DiffinDiffsEstimator,
        DefaultDID,

src/StatsProcedures.jl

@@ -1,5 +1,5 @@
 """
-    StatsStep{Alias,F<:Function}
+    StatsStep{Alias, F<:Function}
 
 Specify the function for moving a step in an [`AbstractStatsProcedure`](@ref).
 An instance of `StatsStep` is callable.
@@ -12,7 +12,7 @@ An instance of `StatsStep` is callable.
     (step::StatsStep{A,F})(ntargs::NamedTuple; verbose::Bool=false)
 
 Call an instance of function of type `F` with arguments
-formed by updating `NamedArgs` with `ntargs`.
+formed by updating `NamedTuple` returned by `[`namedargs(step)`](@ref)` with `ntargs`.
 
 A message with the name of the `StatsStep` is printed to `stdout`
 if a keyword `verbose` takes the value `true`
@@ -21,20 +21,30 @@ The value from `ntargs` supersedes the keyword argument
 in case both are specified.
 
 ## Returns
-- `NamedTuple`: named intermidiate results.
+- `NamedTuple`: named intermediate results.
 """
-struct StatsStep{Alias,F<:Function} end
+struct StatsStep{Alias, F<:Function} end
 
 _f(::StatsStep{A,F}) where {A,F} = F.instance
 
+"""
+    namedargs(s::StatsStep)
+
+Return a `NamedTuple` with keys showing the names of arguments
+accepted by `s` and values representing the defaults.
+"""
+namedargs(s::StatsStep) = error("method for $(typeof(s)) is not defined")
+
 _getargs(ntargs::NamedTuple, s::StatsStep) = _update(ntargs, namedargs(s))
 _update(a::NamedTuple{N1}, b::NamedTuple{N2}) where {N1,N2} =
     NamedTuple{N2}(map(n->getfield(sym_in(n, N1) ? a : b, n), N2))
 
+_combinedargs(::StatsStep, ::Any) = ()
+
 function (step::StatsStep{A,F})(ntargs::NamedTuple; verbose::Bool=false) where {A,F}
     haskey(ntargs, :verbose) && (verbose = ntargs.verbose)
     verbose && printstyled("Running ", step, "\n", color=:green)
-    ret = F.instance(_getargs(ntargs, step)...)
+    ret, share = F.instance(_getargs(ntargs, step)..., _combinedargs(step, (ntargs,))...)
     if ret isa NamedTuple
         return merge(ntargs, ret)
     elseif ret === nothing
@@ -56,15 +66,7 @@ function show(io::IO, ::MIME"text/plain", s::StatsStep{A,F}) where {A,F}
 end
 
 """
-    namedargs(s::StatsStep)
-
-Return a `NamedTuple` with keys showing the names of arguments
-accepted by `s` and values representing the defaults.
-"""
-namedargs(s::StatsStep) = error("method for $(typeof(s)) is not defined")
-
-"""
-    AbstractStatsProcedure{Alias,T<:NTuple{N,StatsStep} where N}
+    AbstractStatsProcedure{Alias, T<:NTuple{N,StatsStep} where N}
 
 Supertype for all types specifying the procedure for statistical estimation or inference.
 
@@ -75,7 +77,7 @@ all subtypes of `AbstractStatsProcedure`.
 - `Alias::Symbol`: alias of the type for pretty-printing.
 - `T<:NTuple{N,StatsStep}`: steps involved in the procedure.
 """
-abstract type AbstractStatsProcedure{A,T<:NTuple{N,StatsStep} where N} end
+abstract type AbstractStatsProcedure{Alias, T<:NTuple{N,StatsStep} where N} end
 
 length(::AbstractStatsProcedure{A,T}) where {A,T} = length(T.parameters)
 eltype(::Type{<:AbstractStatsProcedure}) = StatsStep
@@ -109,7 +111,7 @@ function show(io::IO, ::MIME"text/plain", p::AbstractStatsProcedure{A,T}) where
 end
 
 """
-    SharedStatsStep{T<:StatsStep,I}
+    SharedStatsStep{T<:StatsStep, I}
 
 A [`StatsStep`](@ref) that is possibly shared by
 multiple instances of procedures that are subtypes of [`AbstractStatsProcedure`](@ref).
@@ -119,7 +121,7 @@ See also [`PooledStatsProcedure`](@ref).
 - `T<:StatsStep`: type of the only field `step`.
 - `I`: indices of the procedures that share this step.
 """
-struct SharedStatsStep{T<:StatsStep,I}
+struct SharedStatsStep{T<:StatsStep, I}
     step::T
     function SharedStatsStep(s::StatsStep, pid)
         pid = (unique!(sort!([pid...]))...,)
@@ -130,6 +132,7 @@ end
 _sharedby(::SharedStatsStep{T,I}) where {T,I} = I
 _f(s::SharedStatsStep) = _f(s.step)
 _getargs(ntargs::NamedTuple, s::SharedStatsStep) = _getargs(ntargs, s.step)
+_combinedargs(s::SharedStatsStep, v::AbstractArray) = _combinedargs(s.step, v)
 
 show(io::IO, s::SharedStatsStep) = print(io, s.step)
 
@@ -143,7 +146,7 @@ const SharedStatsSteps = NTuple{N, SharedStatsStep} where N
 const StatsProcedures = NTuple{N, AbstractStatsProcedure} where N
 
 """
-    PooledStatsProcedure{P<:StatsProcedures,S<:SharedStatsSteps}
+    PooledStatsProcedure{P<:StatsProcedures, S<:SharedStatsSteps}
 
 A collection of procedures and shared steps.
 
@@ -155,7 +158,7 @@ See also [`pool`](@ref).
 - `procs::P`: a tuple of instances of subtypes of [`AbstractStatsProcedure`](@ref).
 - `steps::S`: a tuple of [`SharedStatsStep`](@ref) for the procedures in `procs`.
 """
-struct PooledStatsProcedure{P<:StatsProcedures,S<:SharedStatsSteps}
+struct PooledStatsProcedure{P<:StatsProcedures, S<:SharedStatsSteps}
     procs::P
     steps::S
 end
@@ -274,7 +277,7 @@ function show(io::IO, ::MIME"text/plain", ps::PooledStatsProcedure{P,S}) where {
 end
 
 """
-    StatsSpec{Alias,T<:AbstractStatsProcedure}
+    StatsSpec{Alias, T<:AbstractStatsProcedure}
 
 Record the specification for a statistical procedure of type `T`.
 
@@ -297,7 +300,7 @@ or the last value returned by the last [`StatsStep`](@ref) is returned.
 - `keep=nothing`: names (of type `Symbol`) of additional objects to be returned.
 - `keepall::Bool=false`: return all objects returned by each step.
 """
-struct StatsSpec{Alias,T<:AbstractStatsProcedure}
+struct StatsSpec{Alias, T<:AbstractStatsProcedure}
     args::NamedTuple
     StatsSpec(name::Union{Symbol,String},
         T::Type{<:AbstractStatsProcedure}, args::NamedTuple) =
@@ -391,7 +394,7 @@ For end users, `Macro`s that generate `Expr`s for these function calls should be
 
 Optional default arguments are merged
 with the arguments provided for each individual specification
-and replace the default values specified for each procedure.
+and supersede the default values specified for each procedure through [`namedargs`](@ref).
 These default arguments should be specified in the same pattern as
 how arguments are specified for each specification inside the code block,
 as `@specset` processes these arguments by calling
@@ -470,12 +473,18 @@ function proceed(sps::AbstractVector{<:StatsSpec};
         taskids = vcat((gids[steps.procs[i]] for i in _sharedby(step))...)
         tasks = groupview(r->_getargs(r, step), view(traces, taskids))
         for (ins, subtb) in pairs(tasks)
-            ret = _f(step)(ins...)
+            ret, share = _f(step)(ins..., _combinedargs(step, subtb)...)
             ntask += 1
             ntask_total += 1
             if ret !== nothing
-                for i in eachindex(subtb)
-                    subtb[i] = merge(subtb[i], deepcopy(ret))
+                if share
+                    for i in eachindex(subtb)
+                        subtb[i] = merge(subtb[i], ret)
+                    end
+                else
+                    for i in eachindex(subtb)
+                        subtb[i] = merge(subtb[i], deepcopy(ret))
+                    end
                 end
             end
         end

src/parallels.jl

@@ -85,7 +85,7 @@ assume a parallel trends assumption holds over all the relevant time periods.
 abstract type TrendParallel{C,S} <: AbstractParallel{C,S} end
 
 """
-    NeverTreatedParallel{T<:Integer,C,S} <: TrendParallel{C,S}
+    NeverTreatedParallel{C,S,T<:Integer} <: TrendParallel{C,S}
 
 Assume a parallel trends assumption holds between any group
 that received the treatment during the sample periods
@@ -97,15 +97,17 @@ See also [`nevertreated`](@ref).
 - `c::C`: a [`ParallelCondition`](@ref).
 - `s::S`: a [`ParallelStrength`](@ref).
 """
-struct NeverTreatedParallel{T<:Integer,C,S} <: TrendParallel{C,S}
+struct NeverTreatedParallel{C,S,T<:Integer} <: TrendParallel{C,S}
     e::Vector{T}
     c::C
     s::S
     NeverTreatedParallel(e::Vector{T}, c::C, s::S) where
-        {T<:Integer,C<:ParallelCondition,S<:ParallelStrength} =
-            new{T,C,S}(unique!(sort!(e)), c, s)
+        {C<:ParallelCondition,S<:ParallelStrength,T<:Integer} =
+            new{C,S,T}(unique!(sort!(e)), c, s)
 end
 
+treated(pr::NeverTreatedParallel, x) = !(x in pr.e)
+
 show(io::IO, pr::NeverTreatedParallel) =
     print(IOContext(io, :compact=>true), "NeverTreated{", pr.c, ",", pr.s, "}(", pr.e,")")
 
@@ -156,7 +158,7 @@ A wrapper method of `nevertreated` for working with `@formula`.
 @unpack nevertreated
 
 """
-    NotYetTreatedParallel{T<:Integer,C,S} <: TrendParallel{C,S}
+    NotYetTreatedParallel{C,S,T<:Integer} <: TrendParallel{C,S}
 
 Assume a parallel trends assumption holds between any group
 that received the treatment relatively early
@@ -174,17 +176,19 @@ See also [`notyettreated`](@ref).
     - never-treated groups are included and use indices with smaller values;
     - the sample has a rotating panel structure with periods overlapping with some others.
 """
-struct NotYetTreatedParallel{T<:Integer,C,S} <: TrendParallel{C,S}
+struct NotYetTreatedParallel{C,S,T<:Integer} <: TrendParallel{C,S}
     e::Vector{T}
     emin::Union{Vector{T},Nothing}
     c::C
     s::S
     NotYetTreatedParallel(e::Vector{T}, emin::Union{Vector{T},Nothing}, c::C, s::S) where
-        {T<:Integer,C<:ParallelCondition,S<:ParallelStrength} =
-            new{T,C,S}(unique!(sort!(e)),
+        {C<:ParallelCondition,S<:ParallelStrength,T<:Integer} =
+            new{C,S,T}(unique!(sort!(e)),
                 emin isa Nothing ? emin : unique!(sort!(emin)), c, s)
 end
 
+treated(pr::NotYetTreatedParallel, x) = !(x in pr.e)
+
 function show(io::IO, pr::NotYetTreatedParallel)
     print(IOContext(io, :compact=>true), "NotYetTreated{", pr.c, ",", pr.s, "}(", pr.e, ", ")
     print(IOContext(io, :compact=>true), pr.emin isa Nothing ? "NA" : pr.emin, ")")

src/procedures.jl

@@ -1,27 +1,21 @@
 """
     checkdata(args...)
 
-Check `data` is a `Table` and find rows with nonmissing values for variables.
+Check `data` is a `Table` and find valid rows for options `subset` and `weights`.
 See also [`CheckData`](@ref).
-
-# Returns
-- `vars::Vector{Symbol}`: column names of relevant variables.
-- `esample::BitArray`: Boolean indices for rows with nonmissing values for the variables.
 """
-function checkdata(data, tr::AbstractTreatment, pr::AbstractParallel,
-    yterm::AbstractTerm, treatname::Symbol, xterms::TupleTerm,
-    weights::Union{Symbol, Nothing}, subset::Union{AbstractVector, Nothing})
+function checkdata(data, subset::Union{AbstractVector, Nothing},
+        weights::Union{Symbol, Nothing})
 
     istable(data) ||
         throw(ArgumentError("expect `data` being a `Table` while receiving a $(typeof(data))"))
 
-    vars = union([treatname], (termvars(t) for t in (tr, pr, yterm, xterms))...)
-    esample = BitArray(all(v->!ismissing(getproperty(row, v)), vars) for row in rows(data))
-
     if subset !== nothing
         length(subset) != size(data, 1) &&
             throw(DimensionMismatch("`data` of $(size(data, 1)) rows cannot be matched with `subset` vector of $(length(subset)) elements"))
-        esample .&= .!ismissing.(subset) .& subset
+        esample = .!ismissing.(subset) .& subset
+    else
+        esample = trues(size(data, 1))
     end
 
     if weights !== nothing
@@ -30,18 +24,89 @@ function checkdata(data, tr::AbstractTreatment, pr::AbstractParallel,
     end
 
     sum(esample) == 0 && error("no nonmissing data")
-    
-    return (vars=vars, esample=esample)
+    return (esample=esample,), false
 end
 
 """
-    CheckData
+    CheckData <: StatsStep
 
-A [`StatsStep`](@ref) that calls [`DiffinDiffsBase.checkdata`](@ref)
+Call [`DiffinDiffsBase.checkdata`](@ref)
 for some preliminary checks of the input data.
 """
 const CheckData = StatsStep{:CheckData, typeof(checkdata)}
 
-namedargs(::CheckData) = (data=nothing, tr=nothing, pr=nothing, yterm=nothing,
-    treatname=nothing, xterms=nothing, weights=nothing, subset=nothing)
+namedargs(::CheckData) = (data=nothing, subset=nothing, weights=nothing)
+
+function _overlaptime(tr::DynamicTreatment, tr_rows::BitArray, data)
+    control_time = Set(view(getcolumn(data, tr.time), .!tr_rows))
+    treated_time = Set(view(getcolumn(data, tr.time), tr_rows))
+    return intersect(control_time, treated_time), control_time, treated_time
+end
+
+function overlap!(esample::BitArray, tr_rows::BitArray, tr::DynamicTreatment,
+        ::NeverTreatedParallel{Unconditional}, treatname::Symbol, data)
+        overlap_time, control_time, treated_time = _overlaptime(tr, tr_rows, data)
+    length(control_time)==length(treated_time)==length(overlap_time) ||
+        (esample .&= getcolumn(data, tr.time).∈(overlap_time,))
+    tr_rows .&= esample
+end
+    
+function overlap!(esample::BitArray, tr_rows::BitArray, tr::DynamicTreatment,
+        pr::NotYetTreatedParallel{Unconditional}, treatname::Symbol, data)
+    overlap_time, _c, _t = _overlaptime(tr, tr_rows, data)
+    timetype = eltype(overlap_time)
+    if timetype <: Integer
+        emin = pr.emin === nothing ? minimum(pr.e) : pr.emin[1]
+        valid_cohort = filter(x -> x < emin || x in pr.e, overlap_time)
+        filter!(x -> x < emin, overlap_time)
+        esample .&= (getcolumn(data, tr.time).∈(overlap_time,)) .&
+            (getcolumn(data, treatname).∈(valid_cohort,))
+    end
+    tr_rows .&= esample
+end
+
+"""
+    checkvars(args...)
+
+Return rows with observations that are nonmissing and satisfy the overlap condition
+and rows for observations from treated units.
+See also [`CheckVars`](@ref).
+"""
+function checkvars(data, tr::AbstractTreatment, pr::AbstractParallel,
+        yterm::AbstractTerm, treatname::Symbol, treatintterms::TupleTerm,
+        xterms::TupleTerm, esample::BitArray)
+
+    treatvars = union([treatname], (termvars(t) for t in (tr, pr, treatintterms))...)
+    for v in treatvars
+        eltype(getcolumn(data, v)) <: Union{Missing, Integer} ||
+            throw(ArgumentError("data column $v has unaccepted element type"))
+    end
+    # Values of treatintterms from units in control groups are ignored
+    allvars = union(treatvars, (termvars(t) for t in (yterm, xterms))...)
+    treatedvars = setdiff(allvars, termvars(treatintterms))
+    tr_rows = falses(length(esample))
+    @inbounds for i in eachindex(esample)
+        if esample[i]
+            if treated(pr, getcolumn(data, treatname)[i])
+                esample[i] = all(v->!ismissing(getcolumn(data, v)[i]), treatedvars)
+                esample[i] && (tr_rows[i] = true)
+            else
+                esample[i] = all(v->!ismissing(getcolumn(data, v)[i]), allvars)
+            end
+        end
+    end
+
+    overlap!(esample, tr_rows, tr, pr, treatname, data)
+    sum(esample) == 0 && error("no nonmissing data")
+    return (esample=esample, tr_rows=tr_rows), false
+end
+
+"""
+    CheckVars <: StatsStep
+
+Call [`DiffinDiffsBase.checkvars`](@ref) for obtaining valid rows from relevant columns.
+"""
+const CheckVars = StatsStep{:CheckVars, typeof(checkvars)}
 
+namedargs(::CheckVars) = (data=nothing, tr=nothing, pr=nothing,
+    yterm=nothing, treatname=nothing, treatintterms=(), xterms=(), esample=nothing)