Add functions _treatnames and treatnames (#8)

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -8,7 +8,7 @@ using Reexport
 using SplitApplyCombine: groupfind, groupview
 using StatsBase: Weights, uweights
 @reexport using StatsModels
-using Tables: istable, getcolumn, columntable
+using Tables: istable, getcolumn, columntable, columnnames
 
 import Base: ==, show, union
 import Base: eltype, firstindex, lastindex, getindex, iterate, length, sym_in
@@ -68,7 +68,8 @@ export cb,
        didspec,
        @did,
        DIDResult,
-       outcomename
+       outcomename,
+       treatnames
 
 include("utils.jl")
 include("treatments.jl")

src/did.jl

@@ -174,6 +174,27 @@ All concrete subtypes of `DIDResult` are expected to have the following fields:
 """
 abstract type DIDResult <: StatisticalModel end
 
+"""
+    _treatnames(treatinds)
+
+Generate names for treatment coefficients.
+Assume `treatinds` is compatible with the `Tables.jl` interface.
+"""
+function _treatnames(treatinds)
+    cols = columnnames(treatinds)
+    ncol = length(cols)
+    # Assume treatinds has at least one column
+    c1 = cols[1]
+    names = Ref(string(c1, ": ")).*string.(getcolumn(treatinds, c1))
+    if ncol > 1
+        for i in 2:ncol
+            ci = cols[i]
+            names .*= Ref(string(" & ", ci, ": ")).*string.(getcolumn(treatinds, ci))
+        end
+    end
+    return names
+end
+
 """
     coef(r::DIDResult)
 
@@ -297,6 +318,13 @@ Return a vector of coefficient names.
 """
 coefnames(r::DIDResult) = r.coefnames
 
+"""
+    treatnames(r::DIDResult)
+
+Return a vector of names for treatment coefficients.
+"""
+treatnames(r::DIDResult) = r.coefnames[1:size(r.treatinds,1)]
+
 """
     weights(r::DIDResult)
 

test/did.jl

@@ -153,8 +153,15 @@ end
     @test sp6 === @did [noproceed] TestDID @formula(y ~ treat(g, ttreat(t, 0), tpara(0)) & z + x)
 end
 
+@testset "_treatnames" begin
+    t = Table((rel=[1, 2],))
+    @test _treatnames(t) == ["rel: 1", "rel: 2"]
+    r = TestResult(2, 2)
+    @test _treatnames(r.treatinds) == ["rel: $a & c: $b" for a in 1:2 for b in 1:2]
+end
+
 @testset "DIDResult" begin
-    r = result(TestDID, NamedTuple()).result
+    r = TestResult(2, 2)
 
     @test coef(r) == r.coef
     @test coef(r, 1) == 1
@@ -185,6 +192,7 @@ end
     @test responsename(r) == "y"
     @test outcomename(r) == responsename(r)
     @test coefnames(r) == r.coefnames
+    @test treatnames(r) == r.coefnames[1:4]
     @test weights(r) == :w
 end
 

test/runtests.jl

@@ -4,7 +4,8 @@ using DiffinDiffsBase
 using DataFrames
 using DiffinDiffsBase: @fieldequal, unpack, @unpack, hastreat, parse_treat,
     hasintercept, omitsintercept, isintercept, isomitsintercept, parse_intercept,
-    _f, groupargs, pool, checkdata, checkvars!, makeweights, _getsubcolumns, parse_didargs
+    _f, groupargs, pool, checkdata, checkvars!, makeweights, _getsubcolumns, parse_didargs,
+    _treatnames
 using StatsBase: Weights, UnitWeights
 using StatsModels: termvars
 using TypedTables: Table
@@ -26,9 +27,7 @@ const tests = [
 
 printstyled("Running tests:\n", color=:blue, bold=true)
 
-for test in tests
-    @time begin
-        include("$test.jl")
-        println("\033[1m\033[32mPASSED\033[0m: $(test)")
-    end
+@time for test in tests
+    include("$test.jl")
+    println("\033[1m\033[32mPASSED\033[0m: $(test)")
 end