Allow selecting cells and coefficients for contrast (#13)

Author: junyuan-chen

Project.toml

@@ -20,14 +20,14 @@ Vcov = "ec2bfdc2-55df-4fc9-b9ae-4958c2cf2486"
 DataAPI = "1.5"
 DataFrames = "1"
 DiffinDiffsBase = "0.3.3"
-FixedEffectModels = "1.5"
+FixedEffectModels = "1.6"
 FixedEffects = "2"
 Reexport = "0.2, 1"
 StatsBase = "0.33"
 StatsFuns = "0.9"
 StatsModels = "0.6"
 Tables = "1"
-Vcov = "0.4.1"
+Vcov = "0.5"
 julia = "1.3"
 
 [extras]

src/InteractionWeightedDIDs.jl

@@ -20,7 +20,7 @@ using DiffinDiffsBase: ValidTimeType, termvars, isintercept, parse_intercept!,
 import Base: show
 import DiffinDiffsBase: required, default, transformed, combinedargs, copyargs,
     valid_didargs, result, vce, treatment, nobs, outcomename, weights, treatnames,
-    dof_residual, agg, post!, _parse_subset
+    dof_residual, agg, post!
 import FixedEffectModels: has_fe
 
 export Vcov,

src/did.jl

@@ -278,9 +278,10 @@ with weights proportional to `treatweights` within each relative time.
 function agg(r::RegressionBasedDIDResult{<:DynamicTreatment}, names=nothing;
         bys=nothing, subset=nothing)
     inds = subset === nothing ? Colon() : _parse_subset(r, subset, false)
-    ptcells = treatcells(r)
-    bycells = view(ptcells, inds)
-    _parse_bycells!(getfield(bycells, :columns), ptcells, bys)
+    bycells = view(treatcells(r), inds)
+    isempty(bycells) && throw(ArgumentError(
+        "No coefficient left for aggregation after taking the subset"))
+    _parse_bycells!(getfield(bycells, :columns), bycells, bys)
     names === nothing || (bycells = subcolumns(bycells, names, nomissing=false))
 
     tcells, rows = cellrows(bycells, findcell(bycells))
@@ -386,28 +387,43 @@ Base.getproperty(cr::ContrastResult, n::Symbol) = getproperty(_getmat(cr), n)
 Base.parent(cr::ContrastResult) = getfield(cr, :rs)
 
 """
-    contrast(r1::RegDIDResultOrAgg, rs::RegDIDResultOrAgg...)
+    contrast(r1::RegDIDResultOrAgg, rs::RegDIDResultOrAgg...; kwargs)
 
 Construct a [`ContrastResult`](@ref) by collecting the computed least-square weights
 from each of the [`RegDIDResultOrAgg`](@ref).
+
+# Keywords
+- `subset=nothing`: indices for cells to be included (rows in output).
+- `coefs=nothing`: indices for coefficients from each result to be included (columns in output).
 """
-function contrast(r1::RegDIDResultOrAgg, rs::RegDIDResultOrAgg...)
+function contrast(r1::RegDIDResultOrAgg, rs::RegDIDResultOrAgg...;
+        subset=nothing, coefs=nothing)
     has_lsweights(r1) && all(r->has_lsweights(r), rs) || throw(ArgumentError(
         "Results must contain computed least-sqaure weights"))
     ri = r1.lsweights.r
-    ncoef = ntreatcoef(r1)
-    m = r1.lsweights.m
-    for r in rs
-        r.lsweights.r == ri || throw(ArgumentError(
+    rinds = subset === nothing ? Colon() : _parse_subset(ri, subset)
+    # Make a copy to avoid accidentally overwriting cells for DIDResult
+    ri = deepcopy(view(ri, rinds))
+    dcoefs = coefs === nothing ? NamedTuple() : IdDict{Int,Any}(coefs)
+    rs = RegDIDResultOrAgg[r1, rs...]
+    m = Vector{Array}(undef, length(rs)+1)
+    m[1] = view(r1.cellymeans, rinds)
+    iresult = [0]
+    icoef = [0]
+    names = ["cellymeans"]
+    for (i, r) in enumerate(rs)
+        view(r.lsweights.r, rinds) == ri || throw(ArgumentError(
             "Cells for least-square weights comparisons must be identical across the inputs"))
-        ncoef += ntreatcoef(r)
+        cinds = get(dcoefs, i, nothing)
+        cinds = cinds === nothing ? (1:ntreatcoef(r)) : _parse_subset(treatcells(r), cinds)
+        m[i+1] = view(r.lsweights.m, rinds, cinds)
+        ic = view(1:ntreatcoef(r), cinds)
+        push!(iresult, (i for k in 1:length(ic))...)
+        append!(icoef, ic)
+        append!(names, view(treatnames(r), cinds))
     end
-    rs = RegDIDResultOrAgg[r1, rs...]
-    m = hcat(r1.cellymeans, (r.lsweights.m for r in rs)...)
-    rinds = vcat(0, (fill(i+1, ntreatcoef(r)) for (i, r) in enumerate(rs))...)
-    cinds = vcat(0, (1:ntreatcoef(r) for r in rs)...)
-    names = vcat("cellymeans", (treatnames(r) for r in rs)...)
-    ci = VecColumnTable((iresult=rinds, icoef=cinds, name=names))
+    m = hcat(m...)
+    ci = VecColumnTable((iresult=iresult, icoef=icoef, name=names))
     return ContrastResult(rs, TableIndexedMatrix(m, ri, ci))
 end
 
@@ -428,18 +444,8 @@ function Base.sort!(cr::ContrastResult; @nospecialize(kwargs...))
     return cr
 end
 
-_parse_subset(cr::ContrastResult, by::Pair) = (inds = apply(cr.r, by); return inds)
-
-function _parse_subset(cr::ContrastResult, inds)
-    eltype(inds) <: Pair || return inds
-    inds = apply_and(cr.r, inds...)
-    return inds
-end
-
-_parse_subset(::ContrastResult, ::Colon) = Colon()
-
 function Base.view(cr::ContrastResult, subset)
-    inds = _parse_subset(cr, subset)
+    inds = _parse_subset(cr.r, subset)
     r = view(cr.r, inds)
     m = view(cr.m, inds, :)
     return ContrastResult(parent(cr), TableIndexedMatrix(m, r, cr.c))

src/procedures.jl

@@ -508,7 +508,7 @@ function _vce(data, esample::BitVector, vce::Vcov.ClusterCovariance,
     concrete_vce = Vcov.materialize(cludata, vce)
     dof_absorb = 0
     for fe in fes
-        dof_absorb += any(c->isnested(fe, c.refs), concrete_vce.clusters) ? 1 : nunique(fe)
+        dof_absorb += any(c->isnested(fe, c.groups), concrete_vce.clusters) ? 1 : nunique(fe)
     end
     return concrete_vce, dof_absorb
 end

test/did.jl

@@ -212,6 +212,11 @@ end
     @test length(coef(a1)) == 2
     @test coef(a1, "rel: false") ≈ sum(coef(a, "rel: $r") for r in ["-2", "0", "2"])
     @test coef(a1, "rel: true") ≈ sum(coef(a, "rel: $r") for r in ["-3", "1"])
+
+    a2 = agg(r, :rel, bys=:rel=>isodd, subset=:rel=>isodd)
+    @test coef(a2) == coef(a1, 2:2)
+
+    @test_throws ArgumentError agg(r, subset=:rel=>x->x>10)
 end
 
 @testset "@specset" begin
@@ -267,9 +272,12 @@ end
     @test c2.r == r1.lsweights.r == r2.lsweights.r
     @test c2.c.name[1] == "cellymeans"
     @test c2.c.name[2] == r1.coefnames[c2.c.icoef[2]]
+    @test c2.c.iresult == [0, (1 for i in 1:9)..., (2 for i in 1:5)...]
+    @test c2.c.icoef == [0, 1:9..., 1:5...]
     @test parent(c2)[1] === r1
     @test parent(c2)[2] === r2
 
+    @test c1.r !== r1.lsweights.r
     sort!(c1, rev=true)
     @test c1.r == sort(r1.lsweights.r, rev=true)
 
@@ -282,6 +290,25 @@ end
     v = view(c1, [:wave=>x->x==10, :wave_hosp=>x->x==10])
     @test v.m == c1.m[(c1.r.wave.==10).&(c1.r.wave_hosp.==10),:]
     @test v.r == view(c1.r, (c1.r.wave.==10).&(c1.r.wave_hosp.==10))
+
+    c3 = contrast(r1, r2, subset=:wave=>7, coefs=1=>(:wave_hosp,2)=>(x,y)->x+y==8)
+    @test size(c3) == (4, 8)
+    @test parent(c3)[1] === r1
+    @test c3.c.name[2] == "wave_hosp: 8 & rel: 0"
+    @test c3.c.name[3] == "wave_hosp: 10 & rel: -2"
+    @test c3.c.iresult == [0, (1 for i in 1:2)..., (2 for i in 1:5)...]
+    @test c3.c.icoef == [0, 1, 8, 1:5...]
+    @test all(c3.r.wave.==7)
+    ri = r1.lsweights.r.wave.==7
+    @test c3[:,1] == r1.cellymeans[ri]
+    @test c3[:,2:end] == c2[ri,[2,9,11:15...]]
+
+    c4 = contrast(r1, r2, subset=1:3, coefs=(1=>1:2, 2=>:))
+    @test size(c4) == (3, 8)
+    @test c4.c.iresult == [0, (1 for i in 1:2)..., (2 for i in 1:5)...]
+    @test c4.c.icoef == [0:2..., 1:5...]
+    @test c4.r == view(r1.lsweights.r, 1:3)
+    @test c4[:,2:end] == c2[1:3,[2,3,11:15...]]
 end
 
 @testset "post!" begin