Support RotatingTimeArray (#8)

Author: junyuan-chen

src/procedures.jl

@@ -196,7 +196,11 @@ function maketreatcols(data, treatname::Symbol, treatintterms::TermSet,
     cols = subcolumns(data, cellnames, esample)
     cells, rows = cellrows(cols, findcell(cols))
 
-    rel = refarray(cells[2]) .- refarray(cells[1])
+    if cells[1] isa RotatingTimeArray
+        rel = refarray(cells[2].time) .- refarray(cells[1].time)
+    else
+        rel = refarray(cells[2]) .- refarray(cells[1])
+    end
     kept = .!haskey.(Ref(exc), rel) .& .!haskey.(Ref(notreat), cells[1])
     treatrows = rows[kept]
     # Construct cells needed for treatment indicators
@@ -495,7 +499,7 @@ function solveleastsquaresweights(::DynamicTreatment{SharpDesign},
         end
         feM === nothing || _feresiduals!(d, feM, fetol, femaxiter)
         weights isa UnitWeights || (d .*= sqrt.(weights))
-        lswtmat[i,:] .= (crossx \ (X'd))[1:nt]
+        lswtmat[i,:] .= view((crossx \ (X'd)), 1:nt)
     end
     ycellmeans ./= ycellweights
     lswt = TableIndexedMatrix(lswtmat, lswtcells, treatcells)

test/did.jl

@@ -56,16 +56,25 @@
         ──────────────────────────────────────────────────────────────────────"""
 
     df = DataFrame(hrs)
-    df.wave = Date.(df.wave)
-    df.wave_hosp = Date.(df.wave_hosp)
-    df.wave = settime(df, :wave, step=Year(1))
-    df.wave_hosp = settime(df, :wave_hosp, start=Date(7), step=Year(1))
+    df.wave = settime(Date.(hrs.wave), Year(1))
+    df.wave_hosp = settime(Date.(hrs.wave_hosp), Year(1), start=Date(7))
     r1 = @did(Reg, data=df, dynamic(:wave, -1), notyettreated(Date(11)),
         vce=Vcov.cluster(:hhidpn), yterm=term(:oop_spend), treatname=:wave_hosp,
-        treatintterms=(), xterms=(fe(:wave)+fe(:hhidpn)), solvelsweights=false)
+        treatintterms=(), xterms=(fe(:wave)+fe(:hhidpn)), solvelsweights=true)
     @test coef(r1) ≈ coef(r)
     @test r1.coefnames[1] == "wave_hosp: 0008-01-01 & rel: 0"
 
+    rot = ifelse.(isodd.(hrs.hhidpn), 1, 2)
+    df.wave = settime(Date.(hrs.wave), Year(1), rotation=rot)
+    df.wave_hosp = settime(Date.(hrs.wave_hosp), Year(1), start=Date(7), rotation=rot)
+    e = rotatingtime((1,2), Date(11))
+    r2 = @did(Reg, data=df, dynamic(:wave, -1), notyettreated(e),
+        vce=Vcov.cluster(:hhidpn), yterm=term(:oop_spend), treatname=:wave_hosp,
+        treatintterms=(), xterms=(fe(:wave)+fe(:hhidpn)), solvelsweights=true)
+    @test length(coef(r2)) == 18
+    @test coef(r2)[1] ≈ 3790.7218412450593
+    @test r2.coefnames[1] == "wave_hosp: 1_0008-01-01 & rel: 0"
+
     r = @did(Reg, data=hrs, dynamic(:wave, -1), notyettreated([11]),
         vce=Vcov.cluster(:hhidpn), yterm=term(:oop_spend), treatname=:wave_hosp,
         treatintterms=(), cohortinteracted=false, lswtnames=(:wave_hosp, :wave))

test/procedures.jl

@@ -254,10 +254,8 @@ end
     @test combinedargs(MakeTreatCols(), allntargs) ==
         (Dict{Int,Int}(), IdDict{ValidTimeType,Int}())
 
-    df.wave = Date.(df.wave)
-    df.wave_hosp = Date.(df.wave_hosp)
-    df.wave = settime(df, :wave, step=Year(1))
-    df.wave_hosp = settime(df, :wave_hosp, start=Date(7), step=Year(1))
+    df.wave = settime(Date.(hrs.wave), Year(1))
+    df.wave_hosp = settime(Date.(hrs.wave_hosp), Year(1), start=Date(7))
     ret1 = maketreatcols(nt..., typeof(tr), tr.time,
         Dict(-1=>1), IdDict{ValidTimeType,Int}(Date(11)=>1))
     @test ret1.cells[1] == Date.(ret.cells[1])
@@ -270,6 +268,34 @@ end
     @test ret1.cellweights == ret.cellweights
     @test ret1.cellcounts == ret.cellcounts
 
+    rot = ifelse.(isodd.(hrs.hhidpn), 1, 2)
+    df.wave = RotatingTimeArray(rot, hrs.wave)
+    df.wave_hosp = RotatingTimeArray(rot, hrs.wave_hosp)
+    e = rotatingtime((1,1,2), (10,11,11))
+    ret2 = maketreatcols(nt..., typeof(tr), tr.time,
+        Dict(-1=>1), IdDict{ValidTimeType,Int}(c=>1 for c in e))
+    @test ret2.cells[1] == sort!(append!((rotatingtime(r, ret.cells[1]) for r in (1,2))...))
+    rt = append!((rotatingtime(r, 7:11) for r in (1,2))...)
+    @test ret2.cells[2] == repeat(rt, 4)
+    @test size(ret2.treatcells) == (20, 2)
+    @test sort!(unique(ret2.treatcells[1])) ==
+        sort!(append!(rotatingtime(1, 8:9), rotatingtime(2, 8:10)))
+    
+    df.wave = settime(Date.(hrs.wave), Year(1), rotation=rot)
+    df.wave_hosp = settime(Date.(hrs.wave_hosp), Year(1), start=Date(7), rotation=rot)
+    e = rotatingtime((1,1,2), Date.((10,11,11)))
+    ret3 = maketreatcols(nt..., typeof(tr), tr.time,
+        Dict(-1=>1), IdDict{ValidTimeType,Int}(c=>1 for c in e))
+    @test ret3.cells[1].time == Date.(ret2.cells[1].time)
+    @test ret3.cells[2].time == Date.(ret2.cells[2].time)
+    @test ret3.rows == ret2.rows
+    @test ret3.treatcells[1].time == Date.(ret2.treatcells[1].time)
+    @test ret3.treatcells[2] == ret2.treatcells[2]
+    @test ret3.treatrows == ret2.treatrows
+    @test ret3.treatcols == ret2.treatcols
+    @test ret3.cellweights == ret2.cellweights
+    @test ret3.cellcounts == ret2.cellcounts
+
     nt = merge(nt, (data=hrs, tr=tr, pr=pr))
     @test MakeTreatCols()(nt) == merge(nt, (cells=ret.cells, rows=ret.rows,
         treatcells=ret.treatcells, treatrows=ret.treatrows, treatcols=ret.treatcols,