Improve treatment and parallel types and add pause option to proceed (#12)

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -14,6 +14,8 @@ 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
 
+const TimeType = Int
+
 # Reexport objects from StatsBase
 export coef, vcov, responsename, coefnames, weights, nobs, dof_residual
 

src/StatsProcedures.jl

@@ -458,6 +458,7 @@ See also [`@specset`](@ref).
 - `verbose::Bool=false`: print the name of each step when it is called.
 - `keep=nothing`: names (of type `Symbol`) of additional objects to be returned.
 - `keepall::Bool=false`: return all objects generated by procedures along with arguments from the [`StatsSpec`](@ref)s.
+- `pause::Int=0`: break the iteration over [`StatsStep`](@ref)s after finishing the specified number of steps (for debugging).
 
 # Returns
 - `Vector`: results for each specification in the same order of `sps`.
@@ -469,7 +470,7 @@ When either `keep` or `keepall` is specified,
 a `NamedTuple` with additional objects is formed for each [`StatsSpec`](@ref).
 """
 function proceed(sps::AbstractVector{<:StatsSpec};
-        verbose::Bool=false, keep=nothing, keepall::Bool=false)
+        verbose::Bool=false, keep=nothing, keepall::Bool=false, pause::Int=0)
     nsps = length(sps)
     nsps == 0 && throw(ArgumentError("expect a nonempty vector"))
 
@@ -486,6 +487,8 @@ function proceed(sps::AbstractVector{<:StatsSpec};
     steps = pool((p for p in keys(gids))...)
     tasks = IdDict{Tuple, Vector{Int}}()
     ntask_total = 0
+    step_count = 0
+    paused = false
     @inbounds for step in steps
         ntask = 0
         verbose && print("Running ", step, "...")
@@ -521,14 +524,18 @@ function proceed(sps::AbstractVector{<:StatsSpec};
         nprocs = length(_sharedby(step))
         verbose && print("Finished ", ntask, ntask > 1 ? " tasks" : " task", " for ",
             nprocs, nprocs > 1 ? " procedures\n" : " procedure\n")
+        step_count += 1
+        step_count === pause && (paused = true) && break
     end
 
     nprocs = length(steps.procs)
     verbose && printstyled("All steps finished (", ntask_total,
         ntask_total > 1 ? " tasks" : " task", " for ", nprocs,
         nprocs > 1 ? " procedures)\n" : " procedure)\n", bold=true, color=:green)
-    @inbounds for i in 1:nsps
-        traces[i] = result(_procedure(sps[i]), traces[i])
+    if !paused
+        @inbounds for i in 1:nsps
+            traces[i] = result(_procedure(sps[i]), traces[i])
+        end
     end
 
     if keepall
@@ -635,6 +642,7 @@ The following options are available for altering the behavior of `@specset`:
 - `verbose::Bool=false`: print the name of each step when it is called.
 - `keep=nothing`: names (of type `Symbol`) of additional objects to be returned.
 - `keepall::Bool=false`: return all objects generated by procedures along with arguments from the [`StatsSpec`](@ref)s.
+- `pause::Int=0`: break the iteration over [`StatsStep`](@ref)s after finishing the specified number of steps (for debugging).
 """
 macro specset(args...)
     nargs = length(args)

src/did.jl

@@ -134,6 +134,7 @@ The following options are available for altering the behavior of `@did`:
 - `verbose::Bool=false`: print the name of each step when it is called.
 - `keep=nothing`: names (of type `Symbol`) of additional objects to be returned.
 - `keepall::Bool=false`: return all objects generated by procedures along with arguments from the [`StatsSpec`](@ref)s.
+- `pause::Int=0`: break the iteration over [`StatsStep`](@ref)s after finishing the specified number of steps (for debugging).
 """
 macro did(args...)
     nargs = length(args)

src/parallels.jl

@@ -93,16 +93,16 @@ and a group that did not receive any treatment in any sample period.
 See also [`nevertreated`](@ref).
 
 # Fields
-- `e::Vector{Int}`: group indices for units that did not receive any treatment.
+- `e::Tuple{Vararg{TimeType}}`: group indices for units that did not receive any treatment.
 - `c::C`: an instance of [`ParallelCondition`](@ref).
 - `s::S`: an instance of [`ParallelStrength`](@ref).
 """
 struct NeverTreatedParallel{C,S} <: TrendParallel{C,S}
-    e::Vector{Int}
+    e::Tuple{Vararg{TimeType}}
     c::C
     s::S
     function NeverTreatedParallel(e, c::ParallelCondition, s::ParallelStrength)
-        e = unique!(sort!([e...]))
+        e = (unique!(sort!([e...]))...,)
         isempty(e) && error("field `e` cannot be empty") 
         return new{typeof(c),typeof(s)}(e, c, s)
     end
@@ -111,11 +111,12 @@ end
 istreated(pr::NeverTreatedParallel, x) = !(x in pr.e)
 
 show(io::IO, pr::NeverTreatedParallel) =
-    print(IOContext(io, :compact=>true), "NeverTreated{", pr.c, ",", pr.s, "}(", pr.e, ")")
+    print(IOContext(io, :compact=>true), "NeverTreated{", pr.c, ",", pr.s, "}",
+        length(pr.e)==1 ? string("(", pr.e[1], ")") : pr.e)
 
 function show(io::IO, ::MIME"text/plain", pr::NeverTreatedParallel)
     println(io, pr.s, " trends with any never-treated group:")
-    print(io, "  Never-treated groups: ", pr.e)
+    print(io, "  Never-treated groups: ", join(string.(pr.e), ", "))
     pr.c isa Unconditional || print(io, "\n  ", pr.c)
 end
 
@@ -133,14 +134,14 @@ a wrapper method of `nevertreated` calls this method.
 ```jldoctest; setup = :(using DiffinDiffsBase)
 julia> nevertreated(-1)
 Parallel trends with any never-treated group:
-  Never-treated groups: [-1]
+  Never-treated groups: -1
 
 julia> typeof(nevertreated(-1))
-NeverTreatedParallel{Unconditional,Exact,Tuple{Int64}}
+NeverTreatedParallel{Unconditional,Exact}
 
 julia> nevertreated([-1, 0])
 Parallel trends with any never-treated group:
-  Never-treated groups: [-1, 0]
+  Never-treated groups: -1, 0
 
 julia> nevertreated([-1, 0]) == nevertreated(-1:0) == nevertreated(Set([-1, 0]))
 true
@@ -167,8 +168,8 @@ and any group that received the treatment relatively late (or never receved).
 See also [`notyettreated`](@ref).
 
 # Fields
-- `e::Vector{Int}`: group indices for units that received the treatment relatively late.
-- `ecut::Vector{Int}`: user-specified period(s) when units in a group in `e` started to receive treatment.
+- `e::Tuple{Vararg{TimeType}}`: group indices for units that received the treatment relatively late.
+- `ecut::Tuple{Vararg{TimeType}}`: user-specified period(s) when units in a group in `e` started to receive treatment.
 - `c::C`: an instance of [`ParallelCondition`](@ref).
 - `s::S`: an instance of [`ParallelStrength`](@ref).
 
@@ -178,14 +179,14 @@ See also [`notyettreated`](@ref).
     - the sample has a rotating panel structure with periods overlapping with some others.
 """
 struct NotYetTreatedParallel{C,S} <: TrendParallel{C,S}
-    e::Vector{Int}
-    ecut::Vector{Int}
+    e::Tuple{Vararg{TimeType}}
+    ecut::Tuple{Vararg{TimeType}}
     c::C
     s::S
     function NotYetTreatedParallel(e, ecut, c::ParallelCondition, s::ParallelStrength)
-        e = unique!(sort!([e...]))
+        e = (unique!(sort!([e...]))...,)
         isempty(e) && error("field `e` cannot be empty")
-        ecut = unique!(sort!([ecut...]))
+        ecut = (unique!(sort!([ecut...]))...,)
         isempty(ecut) && error("field `ecut` cannot be empty")
         return new{typeof(c),typeof(s)}(e, ecut, c, s)
     end
@@ -194,12 +195,13 @@ end
 istreated(pr::NotYetTreatedParallel, x) = !(x in pr.e)
 
 show(io::IO, pr::NotYetTreatedParallel) =
-    print(IOContext(io, :compact=>true), "NotYetTreated{", pr.c, ",", pr.s, "}(", pr.e, ")")
+    print(IOContext(io, :compact=>true), "NotYetTreated{", pr.c, ",", pr.s, "}",
+        length(pr.e)==1 ? string("(", pr.e[1], ")") : pr.e)
 
 function show(io::IO, ::MIME"text/plain", pr::NotYetTreatedParallel)
     println(io, pr.s, " trends with any not-yet-treated group:")
-    println(io, "  Not-yet-treated groups: ", pr.e)
-    print(io, "  Treated since: ", pr.ecut)
+    println(io, "  Not-yet-treated groups: ", join(string.(pr.e), ", "))
+    print(io, "  Treated since: ", join(string.(pr.ecut), ", "))
     pr.c isa Unconditional || print(io, "\n  ", pr.c)
 end
 
@@ -218,21 +220,21 @@ a wrapper method of `notyettreated` calls this method.
 ```jldoctest; setup = :(using DiffinDiffsBase)
 julia> notyettreated(5)
 Parallel trends with any not-yet-treated group:
-  Not-yet-treated groups: [5]
-  Treated since: [5]
+  Not-yet-treated groups: 5
+  Treated since: 5
 
 julia> typeof(notyettreated(5))
-NotYetTreatedParallel{Unconditional,Exact,Tuple{Int64},Tuple{Int64}}
+NotYetTreatedParallel{Unconditional,Exact}
 
 julia> notyettreated([-1, 5, 6], 5)
 Parallel trends with any not-yet-treated group:
-  Not-yet-treated groups: [-1, 5, 6]
-  Treated since: [5]
+  Not-yet-treated groups: -1, 5, 6
+  Treated since: 5
 
 julia> notyettreated([4, 5, 6], [4, 5, 6])
 Parallel trends with any not-yet-treated group:
-  Not-yet-treated groups: [4, 5, 6]
-  Treated since: [4, 5, 6]
+  Not-yet-treated groups: 4, 5, 6
+  Treated since: 4, 5, 6
 ```
 """
 notyettreated(e, ecut, c::ParallelCondition, s::ParallelStrength) =

src/treatments.jl

@@ -41,27 +41,28 @@ See also [`dynamic`](@ref).
 
 # Fields
 - `time::Symbol`: column name of data representing calendar time.
-- `exc::Vector{Int}`: excluded relative time.
+- `exc::Tuple{Vararg{Int}}`: excluded relative time.
 - `s::S`: an instance of [`TreatmentSharpness`](@ref).
 """
 struct DynamicTreatment{S<:TreatmentSharpness} <: AbstractTreatment
     time::Symbol
-    exc::Vector{Int}
+    exc::Tuple{Vararg{Int}}
     s::S
     function DynamicTreatment(time::Symbol, exc, s::TreatmentSharpness)
-        exc = exc !== nothing ? unique!(sort!([exc...])) : Int[]
+        exc = exc !== nothing ? (unique!(sort!([exc...]))...,) : ()
         return new{typeof(s)}(time, exc, s)
     end
 end
 
 show(io::IO, tr::DynamicTreatment) =
-    print(IOContext(io, :compact=>true), "Dynamic{", tr.s, "}(",
-        isempty(tr.exc) ? "none" : tr.exc, ")")
+    print(IOContext(io, :compact=>true), "Dynamic{", tr.s, "}",
+        length(tr.exc)==1 ? string("(", tr.exc[1], ")") : tr.exc)
 
 function show(io::IO, ::MIME"text/plain", tr::DynamicTreatment)
     println(io, tr.s, " dynamic treatment:")
     println(io, "  column name of time variable: ", tr.time)
-    print(io, "  excluded relative time: ", isempty(tr.exc) ? "none" : tr.exc)
+    print(io, "  excluded relative time: ",
+        isempty(tr.exc) ? "none" : join(string.(tr.exc), ", "))
 end
 
 """
@@ -77,20 +78,20 @@ a wrapper method of `dynamic` calls this method.
 julia> dynamic(:month, -1)
 Sharp dynamic treatment:
   column name of time variable: month
-  excluded relative time: [-1]
+  excluded relative time: -1
 
 julia> typeof(dynamic(:month, -1))
-DynamicTreatment{SharpDesign,Tuple{Int64}}
+DynamicTreatment{SharpDesign}
 
 julia> dynamic(:month, -3:-1)
 Sharp dynamic treatment:
   column name of time variable: month
-  excluded relative time: [-3, -2, -1]
+  excluded relative time: -3, -2, -1
 
 julia> dynamic(:month, [-2,-1], sharp())
 Sharp dynamic treatment:
   column name of time variable: month
-  excluded relative time: [-2, -1]
+  excluded relative time: -2, -1
 ```
 """
 dynamic(time::Symbol, exc, s::TreatmentSharpness=sharp()) =

test/StatsProcedures.jl

@@ -190,6 +190,13 @@ end
     @test p6 == PooledStatsProcedure(ps, shared)
     @test length(p6) == 5
 
+    ps = (rp, cp)
+    shared = [SharedStatsStep(rp[1], 1), SharedStatsStep(rp[2], 1),
+        SharedStatsStep(rp[3], 1), SharedStatsStep(cp[1], 2), SharedStatsStep(cp[2], 2)]
+    p7 = pool(rp, cp)
+    @test p7 == PooledStatsProcedure(ps, shared)
+    @test length(p7) == 5
+
     @test sprint(show, p1) == "PooledStatsProcedure"
     @test sprint(show, MIME("text/plain"), p1) == """
         PooledStatsProcedure with 3 steps from 1 procedure:
@@ -305,6 +312,8 @@ testformatter(nt::NamedTuple) = (haskey(nt, :name) ? nt.name : "", nt.p, (a=nt.a
     @test proceed([s10], keepall=true) == NamedTuple[NamedTuple()]
     @test proceed([s10], keep=:result) == NamedTuple[NamedTuple()]
 
+    @test proceed([s1], pause=1) == ["b"]
+
     @test_throws ArgumentError proceed(StatsSpec[])
 end
 
@@ -354,6 +363,10 @@ end
     r = @specset [verbose keep=[:a]] a=a begin
         StatsSpec(testformatter(testparser(RP; b="b"))...) end
     @test r == [(a="a0", result="a0a0b")]
+
+    r = @specset [verbose pause=1] a=a begin
+        StatsSpec(testformatter(testparser(RP; b="b"))...) end
+    @test r == ["b"]
 
     s0 = @specset [noproceed] for i in 1:3
         a = "a"*string(i)

test/did.jl

@@ -106,20 +106,20 @@ end
         @test sprint(show, sp) == "unnamed"
         @test sprint(show, MIME("text/plain"), sp) == """
             unnamed (StatsSpec for TestDID):
-              Dynamic{S}([-1])
-              NeverTreated{U,P}([-1])"""
+              Dynamic{S}(-1)
+              NeverTreated{U,P}(-1)"""
 
         sp = StatsSpec("", TestDID, (tr=dynamic(:time,-1),))
         @test sprint(show, sp) == "unnamed"
         @test sprint(show, MIME("text/plain"), sp) == """
             unnamed (StatsSpec for TestDID):
-              Dynamic{S}([-1])"""
+              Dynamic{S}(-1)"""
 
-        sp = StatsSpec("name", TestDID, (pr=nevertreated(-1),))
+        sp = StatsSpec("name", TestDID, (pr=notyettreated(-1),))
         @test sprint(show, sp) == "name"
         @test sprint(show, MIME("text/plain"), sp) == """
             name (StatsSpec for TestDID):
-              NeverTreated{U,P}([-1])"""
+              NotYetTreated{U,P}(-1)"""
     end
 end
 

test/parallels.jl

@@ -72,15 +72,15 @@ end
     end
 
     @testset "show" begin
-        @test sprint(show, nt0) == "NeverTreated{U,P}([0])"
+        @test sprint(show, nt0) == "NeverTreated{U,P}(0)"
         @test sprint(show, MIME("text/plain"), nt0) == """
             Parallel trends with any never-treated group:
-              Never-treated groups: [0]"""
+              Never-treated groups: 0"""
 
-        @test sprint(show, nt1) == "NeverTreated{U,P}([0, 1])"
+        @test sprint(show, nt1) == "NeverTreated{U,P}(0, 1)"
         @test sprint(show, MIME("text/plain"), nt1) == """
             Parallel trends with any never-treated group:
-              Never-treated groups: [0, 1]"""
+              Never-treated groups: 0, 1"""
     end
 end
 
@@ -158,29 +158,29 @@ end
     end
 
     @testset "show" begin
-        @test sprint(show, ny0) == "NotYetTreated{U,P}([0])"
+        @test sprint(show, ny0) == "NotYetTreated{U,P}(0)"
         @test sprint(show, MIME("text/plain"), ny0) == """
             Parallel trends with any not-yet-treated group:
-              Not-yet-treated groups: [0]
-              Treated since: [0]"""
+              Not-yet-treated groups: 0
+              Treated since: 0"""
 
-        @test sprint(show, ny1) == "NotYetTreated{U,P}([0, 1])"
+        @test sprint(show, ny1) == "NotYetTreated{U,P}(0, 1)"
         @test sprint(show, MIME("text/plain"), ny1) == """
             Parallel trends with any not-yet-treated group:
-              Not-yet-treated groups: [0, 1]
-              Treated since: [0]"""
+              Not-yet-treated groups: 0, 1
+              Treated since: 0"""
 
-        @test sprint(show, ny2) == "NotYetTreated{U,P}([0, 1])"
+        @test sprint(show, ny2) == "NotYetTreated{U,P}(0, 1)"
         @test sprint(show, MIME("text/plain"), ny2) == """
             Parallel trends with any not-yet-treated group:
-              Not-yet-treated groups: [0, 1]
-              Treated since: [0]"""
+              Not-yet-treated groups: 0, 1
+              Treated since: 0"""
 
-        @test sprint(show, ny3) == "NotYetTreated{U,P}([0, 1])"
+        @test sprint(show, ny3) == "NotYetTreated{U,P}(0, 1)"
         @test sprint(show, MIME("text/plain"), ny3) == """
             Parallel trends with any not-yet-treated group:
-              Not-yet-treated groups: [0, 1]
-              Treated since: [0, 1]"""
+              Not-yet-treated groups: 0, 1
+              Treated since: 0, 1"""
     end
 end