Improve the update of esample and fix issues with overlap! (#28)

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -51,6 +51,7 @@ export cb,
        ScaledVector,
        ScaledMatrix,
        scale,
+       align,
 
        TreatmentSharpness,
        SharpDesign,
@@ -74,6 +75,7 @@ export cb,
        NotYetTreatedParallel,
        notyettreated,
        istreated,
+       istreated!,
 
        TermSet,
        termset,
@@ -84,6 +86,7 @@ export cb,
        findcell,
        cellrows,
        settime,
+       aligntime,
        PanelStructure,
        setpanel,
        findlag!,

src/ScaledArrays.jl

@@ -33,10 +33,19 @@ ScaledArray(rs::RefArray{RA}, pool::P, invpool::Dict{T,R}) where {T,R,RA<:Abstra
 const ScaledVector{T,R} = ScaledArray{T,R,1}
 const ScaledMatrix{T,R} = ScaledArray{T,R,2}
 
-const ScaledArrOrSub = Union{ScaledArray, SubArray{<:Any, <:Any, <:ScaledArray}}
+const ScaledArrOrSub{T,R,N,RA,P} = Union{ScaledArray{T,R,N,RA,P},
+    SubArray{<:Any, <:Any, <:ScaledArray{T,R,N,RA,P}}}
 
+"""
+    scale(sa::ScaledArrOrSub)
+
+Return the step size of the `pool` of `sa`.
+"""
 scale(sa::ScaledArrOrSub) = step(DataAPI.refpool(sa))
 
+Base.size(sa::ScaledArray) = size(sa.refs)
+Base.IndexStyle(::Type{<:ScaledArray{T,R,N,RA}}) where {T,R,N,RA} = IndexStyle(RA)
+
 function _validmin(min, xmin, isstart::Bool)
     if min === nothing
         min = xmin
@@ -180,9 +189,6 @@ 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.size(sa::ScaledArray) = size(sa.refs)
-Base.IndexStyle(::Type{<:ScaledArray{T,R,N,RA}}) where {T,R,N,RA} = IndexStyle(RA)
-
 Base.similar(sa::ScaledArray{T,R}, dims::Dims=size(sa)) where {T,R} =
     ScaledArray(RefArray(ones(R, dims)), DataAPI.refpool(sa), Dict{T,R}())
 
@@ -191,6 +197,28 @@ Base.similar(sa::SubArray{<:Any, <:Any, <:ScaledArray{T,R}}, dims::Dims=size(sa)
 
 Base.similar(sa::ScaledArrOrSub, dims::Int...) = similar(sa, dims)
 
+"""
+    align(xs::AbstractArray, sa::ScaledArrOrSub)
+
+Convert `xs` into a [`ScaledArray`](@ref) with a `pool`
+that has the same first element and step size as the `pool` from `sa`.
+"""
+function align(xs::AbstractArray, sa::ScaledArrOrSub)
+    pool = DataAPI.refpool(sa)
+    invpool = DataAPI.invrefpool(sa)
+    step = scale(sa)
+    xmin, xmax = extrema(xs)
+    start = first(pool)
+    stop = last(pool)
+    start < stop && xmin < start && throw(ArgumentError(
+        "the minimum of xs $xmin is smaller than the minimum of pool $start"))
+    start > stop && xmax > start && throw(ArgumentError(
+        "the maximum of xs $xmax is greater than the maximum of pool $start"))
+    refs = similar(DataAPI.refarray(sa), size(xs))
+    _scaledlabel!(refs, invpool, xs, start, step)
+    return ScaledArray(RefArray(refs), pool, invpool)
+end
+
 DataAPI.refarray(sa::ScaledArray) = sa.refs
 DataAPI.refvalue(sa::ScaledArray, n::Integer) = getindex(DataAPI.refpool(sa), n)
 DataAPI.refpool(sa::ScaledArray) = sa.pool

src/operations.jl

@@ -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) <: ScaledArray || typeof(c) <: SubArray{<:Any,1,<:ScaledArray}
+        if typeof(c) <: ScaledArrOrSub
             columns[i] = similar(c, ncell)
         else
             columns[i] = Vector{eltype(c)}(undef, ncell)
@@ -127,13 +127,15 @@ function cellrows(cols::VecColumnTable, refrows::IdDict)
 end
 
 """
-    settime(data, timename; step, reftype, rotation)
-    settime(time::AbstractArray; step, reftype, rotation)
+    settime(data, timename; step, start, stop, reftype, rotation)
+    settime(time::AbstractArray; step, start, stop, reftype, rotation)
 
-Return a [`ScaledArray`](@ref) that represents discretized time periods.
+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.
 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.
@@ -142,15 +144,18 @@ The returned array ensures well-defined time intervals for operations involving
 
 # 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.
 """
-function settime(time::AbstractArray; step=nothing, reftype::Type{<:Signed}=Int32, rotation=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, step; reftype=reftype)
+    time = ScaledArray(time, start, step, stop; reftype=reftype)
     if rotation !== nothing
         refs = rotatingtime(rotation, time.refs)
         rots = unique(rotation)
@@ -168,10 +173,30 @@ function settime(time::AbstractArray; step=nothing, reftype::Type{<:Signed}=Int3
     return time
 end
 
-function settime(data, timename::Union{Symbol,Integer}; step=nothing,
+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, reftype=reftype, rotation=rotation)
+    return settime(getcolumn(data, timename);
+        step=step, start=start, stop=stop, reftype=reftype, rotation=rotation)
+end
+
+"""
+    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`
+that has the same first element and step size as the `pool` from
+the [`ScaledArray`](@ref) indexed by `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.
+"""
+function aligntime(data, colname::Union{Symbol,Integer}, timename::Union{Symbol,Integer})
+    checktable(data)
+    return align(getcolumn(data, colname), getcolumn(data, timename))
 end
 
 """

src/parallels.jl

@@ -84,6 +84,25 @@ assume a parallel trends assumption holds over all the relevant time periods.
 """
 abstract type TrendParallel{C,S} <: AbstractParallel{C,S} end
 
+"""
+    istreated(pr::TrendParallel, x)
+
+Test whether `x` represents the treatment time
+for a group of units that are not treated.
+See also [`istreated!`](@ref).
+"""
+function istreated end
+
+"""
+    istreated!(out::AbstractVector{Bool}, pr::TrendParallel, x::AbstractArray)
+
+For each element in `x`,
+test whether it represents the treatment time
+for a group of units that are not treated and save the result in `out`.
+See also [`istreated`](@ref).
+"""
+function istreated! end
+
 """
     NeverTreatedParallel{C,S} <: TrendParallel{C,S}
 
@@ -110,6 +129,20 @@ end
 
 istreated(pr::NeverTreatedParallel, x) = !(x in pr.e)
 
+function istreated!(out::AbstractVector{Bool}, pr::NeverTreatedParallel,
+        x::AbstractArray{<:Union{ValidTimeType, Missing}})
+    e = Set(pr.e)
+    out .= .!(x .∈ Ref(e))
+end
+
+function istreated!(out::AbstractVector{Bool}, pr::NeverTreatedParallel,
+        x::ScaledArray{<:Union{ValidTimeType, Missing}})
+    refs = refarray(x)
+    invpool = invrefpool(x)
+    e = Set(invpool[c] for c in pr.e if haskey(invpool, c))
+    out .= .!(refs .∈ Ref(e))
+end
+
 show(io::IO, pr::NeverTreatedParallel) =
     print(IOContext(io, :compact=>true), "NeverTreated{", pr.c, ",", pr.s, "}",
         length(pr.e)==1 ? string("(", pr.e[1], ")") : pr.e)
@@ -196,6 +229,20 @@ end
 
 istreated(pr::NotYetTreatedParallel, x) = !(x in pr.e)
 
+function istreated!(out::AbstractVector{Bool}, pr::NotYetTreatedParallel,
+        x::AbstractArray{<:Union{ValidTimeType, Missing}})
+    e = Set(pr.e)
+    out .= .!(x .∈ Ref(e))
+end
+
+function istreated!(out::AbstractVector{Bool}, pr::NotYetTreatedParallel,
+        x::ScaledArray{<:Union{ValidTimeType, Missing}})
+    refs = refarray(x)
+    invpool = invrefpool(x)
+    e = Set(invpool[c] for c in pr.e if haskey(invpool, c))
+    out .= .!(refs .∈ Ref(e))
+end
+
 show(io::IO, pr::NotYetTreatedParallel) =
     print(IOContext(io, :compact=>true), "NotYetTreated{", pr.c, ",", pr.s, "}",
         length(pr.e)==1 ? string("(", pr.e[1], ")") : pr.e)