Replace Terms with TermSet and allow grouping args by isequal (#14)

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -12,7 +12,7 @@ using Tables: istable, getcolumn, columntable, columnnames
 import Base: ==, show, union
 import Base: eltype, firstindex, lastindex, getindex, iterate, length, sym_in
 import StatsBase: coef, vcov, responsename, coefnames, weights, nobs, dof_residual
-import StatsModels: termvars, hasintercept, omitsintercept
+import StatsModels: termvars
 
 const TimeType = Int
 
@@ -46,7 +46,7 @@ export cb,
        notyettreated,
        istreated,
 
-       Terms,
+       TermSet,
        eachterm,
        TreatmentTerm,
        treat,
@@ -60,6 +60,7 @@ export cb,
        @specset,
 
        CheckData,
+       GroupTerms,
        CheckVars,
        MakeWeights,
 

src/StatsProcedures.jl

@@ -1,12 +1,13 @@
 """
-    StatsStep{Alias, F<:Function}
+    StatsStep{Alias, F<:Function, ById}
 
 Specify the function for moving a step in an [`AbstractStatsProcedure`](@ref).
 An instance of `StatsStep` is callable.
 
 # Parameters
 - `Alias::Symbol`: alias of the type for pretty-printing.
 - `F<:Function`: type of the function to be called by `StatsStep`.
+- `ById::Bool`: whether arguments from multiple [`StatsSpec`](@ref)s should be grouped by `object-id` or `isequal`.
 
 # Methods
     (step::StatsStep{A,F})(ntargs::NamedTuple; verbose::Bool=false)
@@ -23,9 +24,10 @@ in case both are specified.
 ## Returns
 - `NamedTuple`: named intermediate results.
 """
-struct StatsStep{Alias, F<:Function} end
+struct StatsStep{Alias, F<:Function, ById} end
 
 _f(::StatsStep{A,F}) where {A,F} = F.instance
+_byid(::StatsStep{A,F,I}) where {A,F,I} = I
 
 """
     required(s::StatsStep)
@@ -184,6 +186,7 @@ end
 
 _sharedby(s::SharedStatsStep) = s.ids
 _f(s::SharedStatsStep) = _f(s.step)
+_byid(s::SharedStatsStep) = _byid(s.step)
 groupargs(s::SharedStatsStep, @nospecialize(ntargs::NamedTuple)) = groupargs(s.step, ntargs)
 combinedargs(s::SharedStatsStep, v::AbstractArray) = combinedargs(s.step, v)
 copyargs(s::SharedStatsStep) = copyargs(s.step)
@@ -476,17 +479,19 @@ function proceed(sps::Vector{<:StatsSpec};
     end
 
     steps = pool((p for p in keys(gids))...)
-    tasks = IdDict{Tuple, Vector{Int}}()
+    tasks_byid = IdDict{Tuple, Vector{Int}}()
+    tasks_byeq = Dict{Tuple, Vector{Int}}()
     ntask_total = 0
     step_count = 0
     paused = false
     @inbounds for step in steps
+        tasks = _byid(step) ? tasks_byid : tasks_byeq
         ntask = 0
         verbose && print("Running ", step, "...")
-        # Group arguments by ===
+        # Group arguments by objectid or isequal
         for i in _sharedby(step)
-            taskids = gids[steps.procs[i]]
-            for j in taskids
+            traceids = gids[steps.procs[i]]
+            for j in traceids
                 push!(get!(Vector{Int}, tasks, groupargs(step, traces[j])), j)
             end
         end

src/did.jl

@@ -18,13 +18,22 @@ _key(::AbstractTreatment) = :tr
 _key(::AbstractParallel) = :pr
 _key(::Any) = throw(ArgumentError("unacceptable positional arguments"))
 
+function _totermset!(args::Dict{Symbol,Any}, s::Symbol)
+    if haskey(args, s) && !(args[s] isa TermSet)
+        ts = TermSet()
+        foreach(t->setindex!(ts, nothing, t), args[s])
+        args[s] = ts
+    end
+end
+
 """
     parse_didargs!(args::Vector{Any}, kwargs::Dict{Symbol,Any})
 
 Return a `Dict` that is suitable for being passed to
 [`valid_didargs`](@ref) for further processing.
 
 Any [`TreatmentTerm`](@ref) or [`FormulaTerm`](@ref) in `args` is decomposed.
+Any collection of terms is converted to `TermSet`.
 Keys are assigned to all positional arguments based on their types.
 An optional `name` for [`StatsSpec`](@ref) can be included in `args` as a string.
 The order of positional arguments is irrelevant.
@@ -39,8 +48,8 @@ function parse_didargs!(args::Vector{Any}, kwargs::Dict{Symbol,Any})
             kwargs[_key(treat.pr)] = treat.pr
             kwargs[:yterm] = arg.lhs
             kwargs[:treatname] = treat.sym
-            intacts==() || (kwargs[:treatintterms] = intacts)
-            xs==() || (kwargs[:xterms] = xs)
+            kwargs[:treatintterms] = intacts
+            kwargs[:xterms] = xs
         elseif arg isa TreatmentTerm
             kwargs[_key(arg.tr)] = arg.tr
             kwargs[_key(arg.pr)] = arg.pr
@@ -49,6 +58,7 @@ function parse_didargs!(args::Vector{Any}, kwargs::Dict{Symbol,Any})
             kwargs[_key(arg)] = arg
         end
     end
+    foreach(n->_totermset!(kwargs, n), (:treatintterms, :xterms))
     return kwargs
 end
 

src/procedures.jl

@@ -4,7 +4,7 @@
 Check `data` is a `Table` and find valid rows for options `subset` and `weightname`.
 See also [`CheckData`](@ref).
 """
-function checkdata(data, subset::Union{AbstractVector, Nothing},
+function checkdata(data, subset::Union{BitVector, Nothing},
         weightname::Union{Symbol, Nothing})
 
     istable(data) ||
@@ -13,7 +13,7 @@ function checkdata(data, subset::Union{AbstractVector, Nothing},
     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"))
+                cannot be matched with subset vector of $(length(subset)) elements"))
         esample = .!ismissing.(subset) .& subset
     else
         esample = trues(size(data, 1))
@@ -34,11 +34,34 @@ end
 Call [`DiffinDiffsBase.checkdata`](@ref)
 for some preliminary checks of the input data.
 """
-const CheckData = StatsStep{:CheckData, typeof(checkdata)}
+const CheckData = StatsStep{:CheckData, typeof(checkdata), true}
 
 required(::CheckData) = (:data,)
 default(::CheckData) = (subset=nothing, weightname=nothing)
 
+"""
+    groupterms(args...)
+
+Return the arguments for allowing later comparisons based on object-id.
+See also [`GroupTerms`](@ref).
+"""
+groupterms(treatintterms::TermSet, xterms::TermSet) =
+    (treatintterms = treatintterms, xterms = xterms)
+
+"""
+    GroupTerms <: StatsStep
+
+Call [`DiffinDiffsBase.groupterms`](@ref)
+to obtain one of the instances of `treatintterms` and `xterms`
+that have been grouped by `==`
+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}
+
+required(::GroupTerms) = (:treatintterms, :xterms)
+
 function _overlaptime(tr::DynamicTreatment, tr_rows::BitVector, data)
     control_time = Set(view(getcolumn(data, tr.time), .!tr_rows))
     treated_time = Set(view(getcolumn(data, tr.time), tr_rows))
@@ -76,7 +99,7 @@ See also [`CheckVars`](@ref).
 """
 function checkvars!(data, tr::AbstractTreatment, pr::AbstractParallel,
         yterm::AbstractTerm, treatname::Symbol, esample::BitVector,
-        treatintterms::Terms, xterms::Terms)
+        treatintterms::TermSet, xterms::TermSet)
 
     treatvars = union([treatname], (termvars(t) for t in (tr, pr, treatintterms))...)
     for v in treatvars
@@ -105,10 +128,10 @@ end
 
 Call [`DiffinDiffsBase.checkvars!`](@ref) to exclude invalid rows for relevant variables.
 """
-const CheckVars = StatsStep{:CheckVars, typeof(checkvars!)}
+const CheckVars = StatsStep{:CheckVars, typeof(checkvars!), true}
 
 required(::CheckVars) = (:data, :tr, :pr, :yterm, :treatname, :esample)
-default(::CheckVars) = (treatintterms=(), xterms=())
+default(::CheckVars) = (treatintterms=TermSet(), xterms=TermSet())
 copyargs(::CheckVars) = (6,)
 
 """
@@ -133,7 +156,7 @@ end
 
 Call [`DiffinDiffsBase.makeweights`](@ref) to create a generic `Weights` vector.
 """
-const MakeWeights = StatsStep{:MakeWeights, typeof(makeweights)}
+const MakeWeights = StatsStep{:MakeWeights, typeof(makeweights), true}
 
 required(::MakeWeights) = (:data, :esample)
 default(::MakeWeights) = (weightname=nothing,)

src/terms.jl

@@ -1,4 +1,5 @@
 const Terms{N} = NTuple{N, AbstractTerm} where N
+const TermSet = IdDict{AbstractTerm, Nothing}
 
 """
     eachterm(t)
@@ -68,72 +69,95 @@ Extract terms related to treatment specifications from `formula`.
 
 # Returns
 - `TreatmentTerm`: the unique `TreatmentTerm` contained in the `formula`.
-- `Terms`: a tuple of any term that is interacted with the `TreatmentTerm`.
-- `Terms`: a tuple of remaining terms in `formula.rhs`.
+- `TermSet`: a set of terms that are interacted with the `TreatmentTerm`.
+- `TermSet`: a set of remaining terms in `formula.rhs`.
 
 Error will be raised if either existence or uniqueness of the `TreatmentTerm` is violated.
 """
 function parse_treat(@nospecialize(formula::FormulaTerm))
-    # Use Array for detecting duplicate terms
-    treats = Pair{TreatmentTerm,Tuple}[]
+    mettreat = false
+    treatterm = nothing
+    ints = TermSet()
+    xterms = TermSet()
     for term in eachterm(formula.rhs)
-        if hastreat(term)
-            if term isa TreatmentTerm
-                push!(treats, term=>())
-            elseif term isa FunctionTerm
-                push!(treats, treat(term.args_parsed...)=>())
-            elseif term isa InteractionTerm
-                trs = []
-                ints = []
+        if term isa TreatmentTerm
+            if !mettreat
+                mettreat = true
+                treatterm = term
+            else
+                throw(ArgumentError("cannot accept more than one TreatmentTerm"))
+            end
+        elseif term isa FunctionTerm{typeof(treat)}
+            if !mettreat
+                mettreat = true
+                treatterm = treat(term.args_parsed...)
+            else
+                throw(ArgumentError("cannot accept more than one TreatmentTerm"))
+            end
+        elseif term isa InteractionTerm
+            if hastreat(term)
+                !mettreat ||
+                    throw(ArgumentError("cannot accept more than one TreatmentTerm"))
                 for t in term.terms
-                    if hastreat(t)
-                        if t isa TreatmentTerm
-                            push!(trs, t)
-                        elseif t isa FunctionTerm
-                            push!(trs, treat(t.args_parsed...))
+                    if t isa TreatmentTerm
+                        if !mettreat
+                            mettreat = true
+                            treatterm = t
+                        else
+                            throw(ArgumentError("cannot accept more than one TreatmentTerm"))
+                        end
+                    elseif t isa FunctionTerm{typeof(treat)}
+                        if !mettreat
+                            mettreat = true
+                            treatterm = treat(t.args_parsed...)
+                        else
+                            throw(ArgumentError("cannot accept more than one TreatmentTerm"))
                         end
                     else
-                        push!(ints, t)
+                        ints[t] = nothing
                     end
                 end
-                if length(trs)!=1
-                    throw(ArgumentError("invlid term $term in formula.
-                        An interaction term may contain at most one instance of `TreatmentTerm`."))
-                else
-                    push!(treats, trs[1]=>Tuple(ints))
-                end
+            else
+                xterms[term] = nothing
             end
+        else
+            xterms[term] = nothing
         end
     end
-    length(treats)>1 &&
-        throw(ArgumentError("cannot accept more than one `TreatmentTerm`."))
-    isempty(treats) &&
-        throw(ArgumentError("no `TreatmentTerm` is found."))
-    xterms = Tuple(term for term in eachterm(formula.rhs) if !hastreat(term))
-    return treats[1][1], treats[1][2], xterms
+    mettreat || throw(ArgumentError("no TreatmentTerm is found"))
+    return treatterm::TreatmentTerm, ints, xterms
 end
 
-# A tentative solution to changes made in StatsModels v0.6.21
-hasintercept(::Tuple{}) = false
-omitsintercept(::Tuple{}) = false
-
 isintercept(t::AbstractTerm) = t in (InterceptTerm{true}(), ConstantTerm(1))
 isomitsintercept(t::AbstractTerm) =
     t in (InterceptTerm{false}(), ConstantTerm(0), ConstantTerm(-1))
 
 """
-    parse_intercept(ts::Terms)
+    parse_intercept(ts::TermSet)
 
-Convert any `ConstantTerm` to `InterceptTerm` and add them to the end of the tuple.
-This is useful for obtaining a unique way of specifying the intercept
+Convert any `ConstantTerm` to `InterceptTerm`
+and return Boolean values indicating whether terms explictly requiring
+including/excluding the intercept exist.
+
+This function is useful for obtaining a unique way of specifying the intercept
 before going through the `schema`--`apply_schema` pipeline defined in `StatsModels`.
 """
-function parse_intercept(@nospecialize(ts::Terms))
-    out = AbstractTerm[t for t in ts if !(isintercept(t) || isomitsintercept(t))]
-    omitsintercept(ts) && push!(out, InterceptTerm{false}())
-    # This order is assumed by InteractionWeightedDIDs.Fstat
-    hasintercept(ts) && push!(out, InterceptTerm{true}())
-    return (out...,)
+function parse_intercept!(ts::TermSet)
+    hasintercept = false
+    hasomitsintercept = false
+    for t in keys(ts)
+        if isintercept(t)
+            delete!(ts, t)
+            hasintercept = true
+        end
+        if isomitsintercept(t)
+            delete!(ts, t)
+            hasomitsintercept = true
+        end
+    end
+    hasintercept && (ts[InterceptTerm{true}()] = nothing)
+    hasomitsintercept && (ts[InterceptTerm{false}()] = nothing)
+    return hasintercept, hasomitsintercept
 end
 
-termvars(::Tuple{}) = Symbol[]
+termvars(ts::TermSet) = mapreduce(termvars, union, keys(ts), init=Symbol[])

test/StatsProcedures.jl

@@ -3,29 +3,29 @@ using DiffinDiffsBase: _f, _get, groupargs,
 import DiffinDiffsBase: required, default, transformed, combinedargs, copyargs
 
 testvoidstep(a::String) = NamedTuple()
-const TestVoidStep = StatsStep{:TestVoidStep, typeof(testvoidstep)}
+const TestVoidStep = StatsStep{:TestVoidStep, typeof(testvoidstep), true}
 required(::TestVoidStep) = (:a,)
 
 testregstep(a::String, b::String) = (c=a*b,)
-const TestRegStep = StatsStep{:TestRegStep, typeof(testregstep)}
+const TestRegStep = StatsStep{:TestRegStep, typeof(testregstep), true}
 default(::TestRegStep) = (a="a", b="b")
 
 testlaststep(a::String, c::String) = (result=a*c,)
-const TestLastStep = StatsStep{:TestLastStep, typeof(testlaststep)}
+const TestLastStep = StatsStep{:TestLastStep, typeof(testlaststep), true}
 default(::TestLastStep) = (a="a",)
 transformed(::TestLastStep, ntargs::NamedTuple) = (ntargs.c,)
 
 testcombinestep(a::String, bs::String...) = (c=collect(bs),)
-const TestCombineStep = StatsStep{:TestCombineStep, typeof(testcombinestep)}
+const TestCombineStep = StatsStep{:TestCombineStep, typeof(testcombinestep), true}
 default(::TestCombineStep) = (a="a",)
 combinedargs(::TestCombineStep, ntargs) = [nt.b for nt in ntargs]
 
 testarraystep(a::String, c::Array) = (result=c,)
-const TestArrayStep = StatsStep{:TestArrayStep, typeof(testarraystep)}
+const TestArrayStep = StatsStep{:TestArrayStep, typeof(testarraystep), true}
 required(::TestArrayStep) = (:a, :c)
 copyargs(::TestArrayStep) = (2,)
 
-const TestUnnamedStep = StatsStep{:TestUnnamedStep, typeof(testregstep)}
+const TestUnnamedStep = StatsStep{:TestUnnamedStep, typeof(testregstep), true}
 
 @testset "StatsStep" begin
     @testset "_get" begin