Add RotatingTimeArray and improve settime (#30)

Author: junyuan-chen

Project.toml

@@ -18,6 +18,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 StatsModels = "3eaba693-59b7-5ba5-a881-562e759f1c8d"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]
@@ -34,6 +35,7 @@ Reexport = "0.2, 1"
 StatsBase = "0.33"
 StatsFuns = "0.9"
 StatsModels = "0.6.18"
+StructArrays = "0.5"
 Tables = "1.2"
 julia = "1.3"
 

src/DiffinDiffsBase.jl

@@ -17,6 +17,7 @@ using StatsBase: CoefTable, Weights, stderror, uweights
 using StatsFuns: tdistccdf, tdistinvcdf
 @reexport using StatsModels
 using StatsModels: Schema
+using StructArrays: StructArray
 using Tables
 using Tables: AbstractColumns, istable, columnnames, getcolumn
 
@@ -37,7 +38,7 @@ export cb,
 
        RotatingTimeValue,
        rotatingtime,
-       RotatingRange,
+       RotatingTimeArray,
 
        VecColumnTable,
        VecColsRow,

src/ScaledArrays.jl

@@ -8,7 +8,7 @@ end
 """
     ScaledArray{T,R,N,RA,P} <: AbstractArray{T,N}
 
-An array type that stores data as indices of a range.
+Array type that stores data as indices of a range.
 
 # Fields
 - `refs::RA<:AbstractArray{R,N}`: an array of indices.
@@ -189,10 +189,7 @@ ScaledArray(sa::ScaledArray, step=nothing; reftype::Type=eltype(refarray(sa)),
     start=nothing, stop=nothing, xtype::Type=eltype(sa), usepool::Bool=true) =
         ScaledArray(sa, reftype, xtype, start, step, stop, usepool)
 
-Base.similar(sa::ScaledArray{T,R}, dims::Dims=size(sa)) where {T,R} =
-    ScaledArray(RefArray(ones(R, dims)), DataAPI.refpool(sa), Dict{T,R}())
-
-Base.similar(sa::SubArray{<:Any, <:Any, <:ScaledArray{T,R}}, dims::Dims=size(sa)) where {T,R} =
+Base.similar(sa::ScaledArrOrSub{T,R}, dims::Dims=size(sa)) where {T,R} =
     ScaledArray(RefArray(ones(R, dims)), DataAPI.refpool(sa), Dict{T,R}())
 
 Base.similar(sa::ScaledArrOrSub, dims::Int...) = similar(sa, dims)

src/operations.jl

@@ -11,7 +11,7 @@ end
 function _refs_pool(col::AbstractArray, reftype::Type{<:Integer}=UInt32)
     refs = refarray(col)
     pool = refpool(col)
-    labeled = pool !== nothing && !(pool isa RotatingTimeRange)
+    labeled = pool !== nothing
     if !labeled
         refs, invpool, pool = _label(col, eltype(col), reftype)
     end
@@ -97,7 +97,7 @@ function cellrows(cols::VecColumnTable, refrows::IdDict)
     columns = Vector{AbstractVector}(undef, ncol)
     for i in 1:ncol
         c = cols[i]
-        if typeof(c) <: ScaledArrOrSub
+        if typeof(c) <: Union{ScaledArrOrSub, RotatingTimeArray}
             columns[i] = similar(c, ncell)
         else
             columns[i] = Vector{eltype(c)}(undef, ncell)
@@ -127,76 +127,66 @@ function cellrows(cols::VecColumnTable, refrows::IdDict)
 end
 
 """
-    settime(data, timename; step, start, stop, reftype, rotation)
-    settime(time::AbstractArray; step, start, stop, reftype, rotation)
+    settime(time::AbstractArray, step; start, stop, reftype, rotation)
 
 Convert a column of time values to a [`ScaledArray`](@ref)
 for representing discretized time periods of uniform length.
-Time values can be provided either as a table containing the relevant column or as an array.
+If `rotation` is specified (time values belong to multiple rotation groups),
+a [`RotatingTimeArray`](@ref) is returned with the `time` field
+being a [`ScaledArray`](@ref).
 The returned array ensures well-defined time intervals for operations involving relative time
 (such as [`lag`](@ref) and [`diff`](@ref)).
 See also [`aligntime`](@ref).
 
 # Arguments
-- `data`: a Tables.jl-compatible data table.
-- `timename::Union{Symbol,Integer}`: the name of the column in `data` that contains time values.
-- `time::AbstractArray`: the array containing time values (only needed for the alternative method).
+- `time::AbstractArray`: the array containing time values.
+- `step=nothing`: the length of each time interval; try `step=one(eltype(time))` if not specified.
 
 # Keywords
-- `step=nothing`: the length of each time interval; try step=1 if not specified.
 - `start=nothing`: the first element of the `pool` of the returned [`ScaledArray`](@ref).
 - `stop=nothing`: the last element of the `pool` of the returned [`ScaledArray`](@ref).
-- `reftype::Type{<:Signed}=Int32`: the element type of the reference values for the returned [`ScaledArray`](@ref).
-- `rotation=nothing`: rotation groups in a rotating sampling design; use [`RotatingTimeValue`](@ref)s as reference values.
+- `reftype::Type{<:Signed}=Int32`: the element type of the reference values for the [`ScaledArray`](@ref).
+- `rotation=nothing`: rotation groups in a rotating sampling design.
 """
-function settime(time::AbstractArray; step=nothing, start=nothing, stop=nothing,
+function settime(time::AbstractArray, step=nothing; start=nothing, stop=nothing,
         reftype::Type{<:Signed}=Int32, rotation=nothing)
     T = eltype(time)
     T <: ValidTimeType && !(T <: RotatingTimeValue) ||
         throw(ArgumentError("unaccepted element type $T from time column"))
     step === nothing && (step = one(T))
     time = ScaledArray(time, start, step, stop; reftype=reftype)
-    if rotation !== nothing
-        refs = rotatingtime(rotation, time.refs)
-        rots = unique(rotation)
-        invpool = Dict{RotatingTimeValue{eltype(rotation), T}, eltype(refs)}()
-        for (k, v) in time.invpool
-            for r in rots
-                rt = RotatingTimeValue(r, k)
-                invpool[rt] = RotatingTimeValue(r, v)
-            end
-        end
-        rmin, rmax = extrema(rots)
-        pool = RotatingTimeValue(rmin, first(time.pool)):scale(time):RotatingTimeValue(rmax, last(time.pool))
-        time = ScaledArray(RefArray(refs), pool, invpool)
-    end
+    rotation === nothing || (time = RotatingTimeArray(rotation, time))
     return time
 end
 
-function settime(data, timename::Union{Symbol,Integer};
-        step=nothing, start=nothing, stop=nothing,
-        reftype::Type{<:Signed}=Int32, rotation=nothing)
-    checktable(data)
-    return settime(getcolumn(data, timename);
-        step=step, start=start, stop=stop, reftype=reftype, rotation=rotation)
-end
-
 """
+    aligntime(col::AbstractArray, time::ScaledArrOrSub)
+    aligntime(col::AbstractArray, time::RotatingTimeArray)
     aligntime(data, colname::Union{Symbol,Integer}, timename::Union{Symbol,Integer})
 
-Convert a column of time values indexed by `colname` from `data` table
-to a [`ScaledArray`](@ref) with a `pool`
+Convert a column of time values `col` to a [`ScaledArray`](@ref) with a `pool`
 that has the same first element and step size as the `pool` from
-the [`ScaledArray`](@ref) indexed by `timename`.
+the [`ScaledArray`](@ref) `time`.
+If `time` is a [`RotatingTimeArray`](@ref) with the `time` field being a [`ScaledArray`](@ref),
+the returned array is also a [`RotatingTimeArray`](@ref)
+with the `time` field being the converted [`ScaledArray`](@ref).
+Alternative, the arrays may be specified with a Tables.jl-compatible `data` table
+and column indices `colname` and `timename`.
 See also [`settime`](@ref).
 
 This is useful for representing all discretized time periods with the same scale
 so that the underlying reference values returned by `DataAPI.refarray`
 can be directly comparable across the columns.
 """
+aligntime(col::AbstractArray, time::ScaledArrOrSub) = align(col, time)
+aligntime(col::AbstractArray, time::RotatingTimeArray) =
+    RotatingTimeArray(time.rotation, align(col, time.time))
+aligntime(col::RotatingTimeArray, time::RotatingTimeArray) =
+    RotatingTimeArray(time.rotation, align(col.time, time.time))
+
 function aligntime(data, colname::Union{Symbol,Integer}, timename::Union{Symbol,Integer})
     checktable(data)
-    return align(getcolumn(data, colname), getcolumn(data, timename))
+    return aligntime(getcolumn(data, colname), getcolumn(data, timename))
 end
 
 """
@@ -246,7 +236,7 @@ that is returned by [`settime`](@ref).
 - `time::AbstractArray`: the array containing time values (only needed for the alternative method).
 
 # Keywords
-- `step=nothing`: the length of each time interval; try step=1 if not specified.
+- `step=nothing`: the length of each time interval; try `step=one(eltype(time))` if not specified.
 - `reftype::Type{<:Signed}=Int32`: the element type of the reference values for [`PanelStructure`](@ref).
 - `rotation=nothing`: rotation groups in a rotating sampling design; use [`RotatingTimeValue`](@ref)s as reference values.
 
@@ -263,7 +253,7 @@ function setpanel(id::AbstractArray, time::AbstractArray; step=nothing,
         "id has length $(length(id)) while time has length $(length(time))"))
     refs, idpool, labeled = _refs_pool(id)
     labeled && (refs = copy(refs); idpool = copy(idpool))
-    time = settime(time; step=step, reftype=reftype, rotation=rotation)
+    time = settime(time, step; reftype=reftype, rotation=rotation)
     trefs = refarray(time)
     tpool = refpool(time)
     # Multiply 2 to create enough gaps between id groups for the largest possible lead/lag

src/procedures.jl

@@ -67,6 +67,19 @@ const GroupTerms = StatsStep{:GroupTerms, typeof(groupterms), false}
 
 required(::GroupTerms) = (:treatintterms, :xterms)
 
+function _checkscales(col1::AbstractArray, col2::AbstractArray, treatvars::Vector{Symbol})
+    if col1 isa ScaledArrOrSub || col2 isa ScaledArrOrSub
+        col1 isa ScaledArrOrSub && col2 isa ScaledArrOrSub ||
+            throw(ArgumentError("time fields in both columns $(treatvars[1]) and $(treatvars[2]) must be ScaledArrays; see settime and aligntime"))
+        first(DataAPI.refpool(col1)) == first(DataAPI.refpool(col2)) &&
+            scale(col1) == scale(col2) || throw(ArgumentError(
+            "time fields in columns $(treatvars[1]) and $(treatvars[2]) are not aligned; see aligntime"))
+    else
+        eltype(col1) <: Integer || throw(ArgumentError(
+            "columns $(treatvars[1]) and $(treatvars[2]) must be stored in ScaledArrays; see settime and aligntime"))
+    end
+end
+
 function checktreatvars(::DynamicTreatment{SharpDesign}, pr::TrendParallel{Unconditional},
         treatvars::Vector{Symbol}, data)
     # treatvars should be cohort and time variables
@@ -76,22 +89,22 @@ function checktreatvars(::DynamicTreatment{SharpDesign}, pr::TrendParallel{Uncon
     T2 = nonmissingtype(eltype(col2))
     T1 == T2 || throw(ArgumentError(
         "nonmissing elements from columns $(treatvars[1]) and $(treatvars[2]) have different types $T1 and $T2"))
-    T1 <: Union{Integer, RotatingTimeValue{<:Any, <:Integer}} ||
-        col1 isa ScaledArrOrSub && col2 isa ScaledArrOrSub ||
-        throw(ArgumentError("columns $(treatvars[1]) and $(treatvars[2]) must either have integer elements or be ScaledArrays; see settime and aligntime"))
     T1 <: ValidTimeType ||
         throw(ArgumentError("column $(treatvars[1]) has unaccepted element type $(T1)"))
     eltype(pr.e) == T1 || throw(ArgumentError("element type $(eltype(pr.e)) of control cohorts from $pr does not match element type $T1 from data; expect $T1"))
-    if col1 isa ScaledArrOrSub
-        first(DataAPI.refpool(col1)) == first(DataAPI.refpool(col2)) &&
-            scale(col1) == scale(col2) || throw(ArgumentError(
-            "time values in columns $(treatvars[1]) and $(treatvars[2]) are not aligned; see aligntime"))
+    if T1 <: RotatingTimeValue
+        col1 isa RotatingTimeArray && col2 isa RotatingTimeArray ||
+            throw(ArgumentError("columns $(treatvars[1]) and $(treatvars[2]) must be RotatingTimeArrays; see settime"))
+        _checkscales(col1.time, col2.time, treatvars)
+    else
+        _checkscales(col1, col2, treatvars)
     end
 end
 
 function _overlaptime(tr::DynamicTreatment, tr_rows::BitVector, data)
-    control_time = Set(view(refarray(getcolumn(data, tr.time)), .!tr_rows))
-    treated_time = Set(view(refarray(getcolumn(data, tr.time)), tr_rows))
+    timeref = refarray(getcolumn(data, tr.time))
+    control_time = Set(view(timeref, .!tr_rows))
+    treated_time = Set(view(timeref, tr_rows))
     return intersect(control_time, treated_time), control_time, treated_time
 end
 
@@ -108,20 +121,26 @@ end
 function overlap!(esample::BitVector, tr_rows::BitVector, aux::BitVector, tr::DynamicTreatment,
         pr::NotYetTreatedParallel{Unconditional}, treatname::Symbol, data)
     overlap_time, _c, _t = _overlaptime(tr, tr_rows, data)
-    timetype = eltype(overlap_time)
-    invpool = invrefpool(getcolumn(data, tr.time))
-    e = invpool === nothing ? Set(pr.e) : Set(invpool[c] for c in pr.e)
-    if !(timetype <: RotatingTimeValue)
+    timecol = getcolumn(data, tr.time)
+    if !(eltype(timecol) <: RotatingTimeValue)
+        invpool = invrefpool(timecol)
+        e = invpool === nothing ? Set(pr.e) : Set(invpool[c] for c in pr.e)
         ecut = invpool === nothing ? pr.ecut[1] : invpool[pr.ecut[1]]
         filter!(x -> x < ecut, overlap_time)
         isvalidcohort = x -> x < ecut || x in e
     else
-        ecut = invpool === nothing ? pr.ecut : (invpool[e] for e in pr.ecut)
-        ecut = IdDict(e.rotation=>e.time for e in ecut)
+        invpool = invrefpool(timecol.time)
+        if invpool === nothing
+            e = Set(pr.e)
+            ecut = IdDict(e.rotation=>e.time for e in pr.ecut)
+        else
+            e = Set(RotatingTimeValue(c.rotation, invpool[c.time]) for c in pr.e)
+            ecut = IdDict(e.rotation=>invpool[e.time] for e in pr.ecut)
+        end
         filter!(x -> x.time < ecut[x.rotation], overlap_time)
         isvalidcohort = x -> x.time < ecut[x.rotation] || x in e
     end
-    aux[esample] .= view(refarray(getcolumn(data, tr.time)), esample) .∈ (overlap_time,)
+    aux[esample] .= view(refarray(timecol), esample) .∈ (overlap_time,)
     esample[esample] .&= view(aux, esample)
     aux[esample] .= isvalidcohort.(view(refarray(getcolumn(data, treatname)), esample))
     esample[esample] .&= view(aux, esample)