Add more tests and fix some issues

Author: junyuan-chen

src/StatsProcedures.jl

@@ -17,9 +17,11 @@ An instance of `StatsStep` is callable.
 Call an instance of function of type `F` with arguments from `ntargs`
 formed by accessing the keys in `S` and `T` sequentially.
 
-If a keyword argument `verbose` takes `true`
-or `ntargs` contains a key-value pair `verbose=true`,
-a message with the name of the `StatsStep` is printed to `stdout`.
+A message with the name of the `StatsStep` is printed to `stdout`
+if a keyword `verbose` takes the value `true`
+or `ntargs` contains a key-value pair `verbose=true`.
+The value from `ntargs` supersedes the keyword argument
+in case both are specified.
 
 ## Returns
 - `NamedTuple`: named intermidiate results.
@@ -31,8 +33,8 @@ _specnames(::StatsStep{A,F,S}) where {A,F,S} = S
 _tracenames(::StatsStep{A,F,S,T}) where {A,F,S,T} = T
 
 function (step::StatsStep{A,F,S,T})(ntargs::NamedTuple; verbose::Bool=false) where {A,F,S,T}
-    verbose || (haskey(ntargs, :verbose) && ntargs.verbose) &&
-        println("  Running ", step)
+    haskey(ntargs, :verbose) && (verbose = ntargs.verbose)
+    verbose && printstyled("Running ", step, "\n", color=:green)
     args = NamedTuple{(S...,T...)}(ntargs)
     ret = F.instance(args...)
     if ret isa NamedTuple
@@ -323,7 +325,7 @@ while ignoring the orders.
 function (sp::StatsSpec{A,T})(;
         verbose::Bool=false, keep=nothing, keepall::Bool=false) where {A,T}
     args = verbose ? merge(sp.args, (verbose=true,)) : sp.args
-    ntall = foldl(|>, T(), init=sp.args)
+    ntall = foldl(|>, T(), init=args)
     if keepall
         return ntall
     elseif !isempty(ntall)
@@ -336,7 +338,7 @@ function (sp::StatsSpec{A,T})(;
             elseif eltype(keep) != Symbol
                 throw(ArgumentError("expect Symbol or collections of Symbols for the value of option `keep`"))
             end
-            return (result=res, (kv for kv in pairs(ntall) if kv[1] in keep)...)
+            return ((kv for kv in pairs(ntall) if kv[1] in keep)..., result=res)
         end
     else
         return nothing
@@ -361,7 +363,7 @@ function run_specset(sps::AbstractVector{<:StatsSpec};
     gids = groupfind(r->procedure(r.spec), tb)
     steps = pool((p() for p in keys(gids))...)
     for step in steps
-        verbose && println("  Running ", step)
+        verbose && printstyled("Running ", step, "\n", color=:green)
         args = _specnames(step)
         tras = _tracenames(step)
         byf = r->merge(NamedTuple{args}(r.spec.args), NamedTuple{tras}(r.trace))

src/parallels.jl

@@ -243,6 +243,13 @@ A wrapper method of `notyettreated` for working with `@formula`.
 """
 @unpack notyettreated
 
-termvars(c::ParallelCondition) = Symbol[]
-termvars(s::ParallelStrength) = Symbol[]
-termvars(pr::AbstractParallel) = union(termvars(pr.c), termvars(pr.s))
+termvars(c::ParallelCondition) =
+    error("StatsModels.termvars is not defined for $(typeof(c))")
+termvars(::Unconditional) = Symbol[]
+termvars(s::ParallelStrength) =
+    error("StatsModels.termvars is not defined for $(typeof(s))")
+termvars(::Exact) = Symbol[]
+termvars(pr::AbstractParallel) =
+    error("StatsModels.termvars is not defined for $(typeof(pr))")
+termvars(pr::NeverTreatedParallel) = union(termvars(pr.c), termvars(pr.s))
+termvars(pr::NotYetTreatedParallel) = union(termvars(pr.c), termvars(pr.s))

src/procedures.jl

@@ -1,31 +1,45 @@
+"""
+    checkdata(args...)
 
-function check_data(data, tr::AbstractTreatment, pr::AbstractParallel,
+Check `data` is a `Table` and find rows with nonmissing values for variables.
+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})
 
     istable(data) ||
         throw(ArgumentError("expected data in a Table, got $(typeof(data))"))
 
-    vars = union(treatname, (termvars(t) for t in (tr, pr, yterm, xterms)))
+    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
+    if subset !== nothing
         length(subset) != size(data, 1) &&
-            throw("df has $(size(df, 1)) rows but the subset vector has $(length(subset)) elements")
-        esample .&= .!ismissing.(x) .& x
+            throw(DimensionMismatch("`data` of $(size(data, 1)) rows cannot be matched with `subset` vector of $(length(subset)) elements"))
+        esample .&= .!ismissing.(subset) .& subset
     end
 
-    if weights != nothing
+    if weights !== nothing
         colweights = getcolumn(columns(data), weights)
         esample .&= .!ismissing.(colweights) .& (colweights .> 0)
     end
 
-    sum(esample) == 0 && throw(ArgumentError("no nonmissing data"))
+    sum(esample) == 0 && error("no nonmissing data")
 
-    return (vars=vars, esample=esample,)
+    return (vars=vars, esample=esample)
 end
 
-const CheckData = StatsStep{:CheckData, typeof(check_data), (:data, :tr, :pr, :yterm, :treatname, :xterms, :weights, :subset), ()}
+"""
+    CheckData
 
+A [`StatsStep`](@ref) that calls [`DiffinDiffsBase.checkdata`](@ref)
+for some preliminary checks of the input data.
+"""
+const CheckData = StatsStep{:CheckData, typeof(checkdata), (:data, :tr, :pr, :yterm, :treatname, :xterms, :weights, :subset), ()}
 
 

src/treatments.jl

@@ -41,7 +41,7 @@ abstract type AbstractTreatment end
 @fieldequal AbstractTreatment
 
 """
-    DynamicTreatment{E<:EleOrVec{<:Integer},S<:TreatmentSharpness} <: AbstractTreatment
+    DynamicTreatment{E<:EleOrVec{<:Integer}, S<:TreatmentSharpness} <: AbstractTreatment
 
 Specify an absorbing binary treatment with effects allowed to evolve over time.
 See also [`dynamic`](@ref).
@@ -51,18 +51,14 @@ See also [`dynamic`](@ref).
 - `exc::E`: excluded relative time (either an integer or vector of integers).
 - `s::S`: a [`TreatmentSharpness`](@ref).
 """
-struct DynamicTreatment{E<:EleOrVec{<:Integer},S<:TreatmentSharpness} <: AbstractTreatment
+struct DynamicTreatment{E<:EleOrVec{<:Integer}, S<:TreatmentSharpness} <: AbstractTreatment
     time::Symbol
     exc::E
     s::S
-    function DynamicTreatment(time::Symbol, exc::E, s::S) where
-        {E<:EleOrVec{<:Integer},S<:TreatmentSharpness}
-        if length(exc) > 1
-            exc = sort!(exc)
-        elseif E <: Vector
-            exc = exc[1]
-        end
-        return new{typeof(exc),S}(time, exc, s)
+    function DynamicTreatment(time::Symbol, exc, s::TreatmentSharpness)
+        exc = unique!(sort!([exc...]))
+        length(exc)==1 && (exc = exc[1])
+        return new{typeof(exc),typeof(s)}(time, exc, s)
     end
 end
 
@@ -105,7 +101,7 @@ Sharp dynamic treatment:
 ```
 """
 dynamic(time::Symbol, exc, s::TreatmentSharpness=sharp()) =
-    DynamicTreatment(time, [exc...], s)
+    DynamicTreatment(time, exc, s)
 
 """
     dynamic(ts::AbstractTerm...)
@@ -114,6 +110,9 @@ A wrapper method of `dynamic` for working with `@formula`.
 """
 @unpack dynamic
 
-termvars(s::TreatmentSharpness) = Symbol[]
-termvars(tr::AbstractTreatment) = termvars(tr.s)
-termvars(tr::DynamicTreatment) = pushfirst!(termvars(s::TreatmentSharpness), tr.time)
+termvars(s::TreatmentSharpness) =
+    error("StatsModels.termvars is not defined for $(typeof(s))")
+termvars(::SharpDesign) = Symbol[]
+termvars(tr::AbstractTreatment) =
+    error("StatsModels.termvars is not defined for $(typeof(tr))")
+termvars(tr::DynamicTreatment) = [tr.time, termvars(tr.s)...]

test/StatsProcedures.jl

@@ -1,7 +1,7 @@
 using DiffinDiffsBase: _f, _specnames, _tracenames,
-    _sharedby
+    _sharedby, _show_args
 
-testvoidstep(a::String, b::String) = println(a)
+testvoidstep(a::String, b::String) = nothing
 const TestVoidStep = StatsStep{:TestVoidStep, typeof(testvoidstep), (:a, :b), ()}
 
 testregstep(a::String, b::String) = (b=a, c=a*b,)
@@ -192,3 +192,32 @@ end
           NullProcedure"""
 end
 
+@testset "StatsSpec" begin
+    s1 = StatsSpec("name", RP, (a="a",b="b"))
+    s2 = StatsSpec("", RP, (a="a",b="b"))
+    s3 = StatsSpec("", UP, (a="a",b="b"))
+    s4 = StatsSpec("name", RP, (b="b", a="a"))
+    s5 = StatsSpec("name", RP, (b="b", a="a", d="d"))
+    @test s1 == s2
+    @test s2 != s3
+    @test s2 != s4
+    @test s2 ≊ s4
+
+    @test s1() == "aab"
+    @test s3() == "ab"
+    
+    @test s1(keep=:a) == (a="a", result="aab")
+    @test s1(keep=(:a,:c)) == (a="a", c="ab", result="aab")
+    @test_throws ArgumentError s1(keep=1)
+    @test s1(keepall=true) == (a="a", b="a", c="ab", result="aab")
+
+    s6 = StatsSpec("", NP, NamedTuple())
+    @test s6() === nothing
+
+    @test sprint(show, s1) == "name"
+    @test sprint(show, s2) == "unnamed"
+    @test sprint(show, MIME("text/plain"), s1) == "name (StatsSpec for RegProcedure)"
+    @test sprint(show, MIME("text/plain"), s2) == "unnamed (StatsSpec for RegProcedure)"
+
+    @test _show_args(stdout, s1) === nothing
+end

test/parallels.jl

@@ -170,3 +170,14 @@ end
               Treated since: [0, 1]"""
     end
 end
+
+@testset "termvars" begin
+    @test termvars(Unconditional()) == Symbol[]
+    @test termvars(Exact()) == Symbol[]
+    @test termvars(nevertreated(-1)) == Symbol[]
+    @test termvars(notyettreated(5)) == Symbol[]
+
+    @test_throws ErrorException termvars(TestParaCondition())
+    @test_throws ErrorException termvars(TestParaStrength())
+    @test_throws ErrorException termvars(PR)
+end

test/procedures.jl

@@ -1,16 +1,37 @@
-using DiffinDiffsBase: _f, _specnames, _tracenames,
-    check_data
+using DiffinDiffsBase: _f, _specnames, _tracenames, checkdata
 
-@testset "check_data" begin
+@testset "CheckData" begin
     @testset "StatsStep" begin
         @test sprint(show, CheckData()) == "CheckData"
         @test sprint(show, MIME("text/plain"), CheckData()) == """
-        CheckData (StatsStep that calls DiffinDiffsBase.check_data):
+        CheckData (StatsStep that calls DiffinDiffsBase.checkdata):
           arguments from StatsSpec: (:data, :tr, :pr, :yterm, :treatname, :xterms, :weights, :subset)
           arguments from trace: ()"""
 
-        @test _f(CheckData()) == check_data
+        @test _f(CheckData()) == checkdata
         @test _specnames(CheckData()) == (:data, :tr, :pr, :yterm, :treatname, :xterms, :weights, :subset)
         @test _tracenames(CheckData()) == ()
     end
+
+    @testset "checkdata" begin
+        hrs = exampledata("hrs")
+        nt = (data=hrs, tr=dynamic(:wave, -1), pr=nevertreated(11),
+            yterm=term(:oop_spend), treatname=:wave_hosp, xterms=(),
+            weights=nothing, subset=nothing)
+        @test checkdata(nt...) ==
+            (vars=[:wave_hosp, :wave, :oop_spend], esample=trues(size(hrs,1)))
+        
+        nt = merge(nt, (weights=:rwthh, subset=hrs.male))
+        @test checkdata(nt...) ==
+            (vars=[:wave_hosp, :wave, :oop_spend], esample=BitArray(hrs.male))
+        
+        nt = merge(nt, (data=rand(10,10),))
+        @test_throws ArgumentError checkdata(nt...)
+
+        nt = merge(nt, (data=hrs, subset=BitArray(hrs.male[1:100])))
+        @test_throws DimensionMismatch checkdata(nt...)
+
+        nt = merge(nt, (subset=falses(size(hrs,1)),))
+        @test_throws ErrorException checkdata(nt...)
+    end
 end 

test/runtests.jl

@@ -1,6 +1,8 @@
 using Test
 using DiffinDiffsBase
 
+using StatsModels: termvars
+
 include("testutils.jl")
 
 const tests = [

test/testutils.jl

@@ -4,6 +4,10 @@ import Base: show
 
 sprintcompact(x) = sprint(show, x; context=:compact=>true)
 
+struct TestSharpness <: TreatmentSharpness end
+struct TestParaCondition <: ParallelCondition end
+struct TestParaStrength <: ParallelStrength end
+
 struct TestTreatment <: AbstractTreatment
     time::Symbol
     ref::Int

test/treatments.jl

@@ -56,3 +56,11 @@ end
               excluded relative time: [-2, -1]"""
     end
 end
+
+@testset "termvars" begin
+    @test termvars(sharp()) == Symbol[]
+    @test termvars(dynamic(:month, -1)) == [:month]
+
+    @test_throws ErrorException termvars(TestSharpness())
+    @test_throws ErrorException termvars(TR)
+end