Improve ScaledArray and add RotatingTimeRange (#24)

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -32,8 +32,10 @@ export coef, vcov, stderror, confint, nobs, dof_residual, responsename, coefname
 export cb,
        ≊,
        exampledata,
+
        RotatingTimeValue,
        rotatingtime,
+       RotatingRange,
 
        VecColumnTable,
        VecColsRow,
@@ -46,6 +48,7 @@ export cb,
        ScaledArray,
        ScaledVector,
        ScaledMatrix,
+       scale,
 
        TreatmentSharpness,
        SharpDesign,
@@ -126,6 +129,7 @@ export cb,
        rescale
 
 include("utils.jl")
+include("time.jl")
 include("tables.jl")
 include("ScaledArrays.jl")
 include("treatments.jl")

src/ScaledArrays.jl

@@ -6,34 +6,31 @@ mutable struct RefArray{R}
 end
 
 """
-    ScaledArray{T,N,RA,S,P} <: AbstractArray{T,N}
+    ScaledArray{T,R,N,RA,P} <: AbstractArray{T,N}
 
 An array type that stores data as indices of a range.
 
 # Fields
 - `refs::RA<:AbstractArray{<:Any, N}`: an array of indices.
-- `start::T`: the starting value of the range.
-- `step::S`: the step size of the range.
-- `stop::T`: the stopping value of the range.
 - `pool::P<:AbstractRange{T}`: a range that covers all possible values stored by the array.
+- `invpool::Dict{T,R}`: a map from array elements to indices of `pool`.
 """
-mutable struct ScaledArray{T,N,RA,S,P} <: AbstractArray{T,N}
+mutable struct ScaledArray{T,R,N,RA,P} <: AbstractArray{T,N}
     refs::RA
-    start::T
-    step::S
-    stop::T
     pool::P
-    ScaledArray{T,N,RA,S,P}(rs::RefArray{RA}, start::T, step::S, stop::T,
-        pool::P=start:step:stop) where
-        {T, N, RA<:AbstractArray{<:Any, N}, S, P<:AbstractRange{T}} =
-            new{T,N,RA,S,P}(rs.a, start, step, stop, pool)
+    invpool::Dict{T,R}
+    ScaledArray{T,R,N,RA,P}(rs::RefArray{RA}, pool::P, invpool::Dict{T,R}) where
+        {T, R, N, RA<:AbstractArray{R, N}, P<:AbstractRange{T}} =
+            new{T,R,N,RA,P}(rs.a, pool, invpool)
 end
 
-ScaledArray(rs::RefArray{RA}, start::T, step::S, stop::T, pool::P=start:step:stop) where
-    {T,RA,S,P} = ScaledArray{T,ndims(RA),RA,S,P}(rs, start, step, stop, pool)
+ScaledArray(rs::RefArray{RA}, pool::P, invpool::Dict{T,R}) where
+    {T,R,RA<:AbstractArray{R},P} = ScaledArray{T,R,ndims(RA),RA,P}(rs, pool, invpool)
 
-const ScaledVector{T} = ScaledArray{T,1}
-const ScaledMatrix{T} = ScaledArray{T,2}
+const ScaledVector{T,R} = ScaledArray{T,R,1}
+const ScaledMatrix{T,R} = ScaledArray{T,R,2}
+
+scale(sa::ScaledArray) = step(sa.pool)
 
 function _validmin(min, xmin, isstart::Bool)
     if min === nothing
@@ -55,10 +52,11 @@ function _validmax(max, xmax, isstart::Bool)
     return max
 end
 
-function validstartstepstop(x::AbstractArray, start, step, stop, usepool)
+function validpool(x::AbstractArray, T::Type, start, step, stop, usepool::Bool)
     step === nothing && throw(ArgumentError("step cannot be nothing"))
     pool = DataAPI.refpool(x)
-    xmin, xmax = usepool && pool !== nothing ? extrema(pool) : extrema(x)
+    xs = skipmissing(usepool && pool !== nothing ? pool : x)
+    xmin, xmax = extrema(xs)
     applicable(+, xmin, step) || throw(ArgumentError(
         "step of type $(typeof(step)) does not match array with element type $(eltype(x))"))
     if xmin + step > xmin
@@ -70,47 +68,83 @@ function validstartstepstop(x::AbstractArray, start, step, stop, usepool)
     else
         throw(ArgumentError("step cannot be zero"))
     end
-    T = promote_type(eltype(x), eltype(start:step:stop))
-    return convert(T, start), convert(T, stop)
+    start = convert(T, start)
+    stop = convert(T, stop)
+    return start:step:stop
 end
 
-function _scaledlabel(x::AbstractArray, step, reftype::Type{<:Signed}=DEFAULT_REF_TYPE;
-        start=nothing, stop=nothing, usepool::Bool=true)
-    start, stop = validstartstepstop(x, start, step, stop, usepool)
-    pool = start:step:stop
-    while typemax(reftype) < length(pool)
-        reftype = widen(reftype)
+function _scaledlabel!(labels::AbstractArray, invpool::Dict, xs::AbstractArray, start, step)
+    z = zero(valtype(invpool))
+    @inbounds for i in eachindex(labels)
+        x = xs[i]
+        lbl = get(invpool, x, z)
+        if lbl !== z
+            labels[i] = lbl
+        elseif ismissing(x)
+            labels[i] = z
+            invpool[x] = z
+        else
+            r = start:step:x
+            lbl = length(r)
+            labels[i] = lbl
+            invpool[x] = lbl
+        end
     end
-    refs = similar(x, reftype)
-    @inbounds for i in eachindex(refs)
-        refs[i] = length(start:step:x[i])
+end
+
+function scaledlabel(xs::AbstractArray, stepsize,
+        R::Type=DEFAULT_REF_TYPE, T::Type=eltype(xs);
+        start=nothing, stop=nothing, usepool::Bool=true)
+    pool = validpool(xs, T, start, stepsize, stop, usepool)
+    T = Missing <: T ? Union{eltype(pool), Missing} : eltype(pool)
+    start = first(pool)
+    stepsize = step(pool)
+    if R <: Integer
+        while typemax(R) < length(pool)
+            R = widen(R)
+        end
     end
-    return refs, start, step, stop
+    labels = similar(xs, R)
+    invpool = Dict{T,R}()
+    _scaledlabel!(labels, invpool, xs, start, stepsize)
+    return labels, pool, invpool
 end
 
-function ScaledArray(x::AbstractArray, reftype::Type, start, step, stop, usepool::Bool=true)
-    refs, start, step, stop = _scaledlabel(x, step, reftype; start=start, stop=stop, usepool=usepool)
-    return ScaledArray(RefArray(refs), start, step, stop)
+function ScaledArray(x::AbstractArray, reftype::Type, xtype::Type, start, step, stop, usepool::Bool=true)
+    refs, pool, invpool = scaledlabel(x, step, reftype, xtype; start=start, stop=stop, usepool=usepool)
+    return ScaledArray(RefArray(refs), pool, invpool)
 end
 
-function ScaledArray(sa::ScaledArray, reftype::Type, start, step, stop, usepool::Bool=true)
-    if step !== nothing && step != sa.step
-        refs, start, step, stop = _scaledlabel(sa, step, reftype; start=start, stop=stop, usepool=usepool)
-        return ScaledArray(RefArray(refs), start, step, stop)
+function ScaledArray(sa::ScaledArray, reftype::Type, xtype::Type, start, step, stop, usepool::Bool=true)
+    if step !== nothing && step != scale(sa)
+        refs, pool, invpool = scaledlabel(sa, step, reftype, xtype; start=start, stop=stop, usepool=usepool)
+        return ScaledArray(RefArray(refs), pool, invpool)
     else
-        step = sa.step
-        start, stop = validstartstepstop(sa, start, step, stop, usepool)
+        step = scale(sa)
+        pool = validpool(sa, xtype, start, step, stop, usepool)
+        T = Missing <: xtype ? Union{eltype(pool), Missing} : eltype(pool)
         refs = similar(sa.refs, reftype)
-        if start == sa.start
+        invpool = Dict{T, reftype}()
+        start0 = first(sa.pool)
+        start = first(pool)
+        stop = last(pool)
+        if start == start0
             copy!(refs, sa.refs)
-        elseif start < sa.start && start < stop || start > sa.start && start > stop
-            offset = length(start:step:sa.start) - 1
+            copy!(invpool, sa.invpool)
+        elseif start < start0 && start < stop || start > start0 && start > stop
+            offset = length(start:step:start0) - 1
             refs .= sa.refs .+ offset
+            for (k, v) in sa.invpool
+                invpool[k] = v + offset
+            end
         else
-            offset = length(sa.start:step:start) - 1
+            offset = length(start0:step:start) - 1
             refs .= sa.refs .- offset
+            for (k, v) in sa.invpool
+                invpool[k] = v - offset
+            end
         end
-        return ScaledArray(RefArray(refs), start, step, stop)
+        return ScaledArray(RefArray(refs), pool, invpool)
     end
 end
 
@@ -126,39 +160,43 @@ If `start` or `stop` is not specified, it will be chosen based on the extrema of
 - `usepool::Bool=true`: find extrema of `x` based on `DataAPI.refpool`.
 """
 ScaledArray(x::AbstractArray, start, step, stop=nothing;
-    reftype::Type=DEFAULT_REF_TYPE, usepool::Bool=true) =
-        ScaledArray(x, reftype, start, step, stop, usepool)
+    reftype::Type=DEFAULT_REF_TYPE, xtype::Type=eltype(x), usepool::Bool=true) =
+        ScaledArray(x, reftype, xtype, start, step, stop, usepool)
 
 ScaledArray(sa::ScaledArray, start, step, stop=nothing;
-    reftype::Type=eltype(refarray(sa)), usepool::Bool=true) =
-        ScaledArray(sa, reftype, start, step, stop, usepool)
+    reftype::Type=eltype(refarray(sa)), xtype::Type=eltype(sa), usepool::Bool=true) =
+        ScaledArray(sa, reftype, xtype, start, step, stop, usepool)
 
 ScaledArray(x::AbstractArray, step; reftype::Type=DEFAULT_REF_TYPE,
-    start=nothing, stop=nothing, usepool::Bool=true) =
-        ScaledArray(x, reftype, start, step, stop, usepool)
+    start=nothing, stop=nothing, xtype::Type=eltype(x), usepool::Bool=true) =
+        ScaledArray(x, reftype, xtype, start, step, stop, usepool)
 
 ScaledArray(sa::ScaledArray, step=nothing; reftype::Type=eltype(refarray(sa)),
-    start=nothing, stop=nothing, usepool::Bool=true) =
-        ScaledArray(sa, reftype, start, step, stop, usepool)
+    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,N,RA}}) where {T,N,RA} = IndexStyle(RA)
+Base.IndexStyle(::Type{<:ScaledArray{T,R,N,RA}}) where {T,R,N,RA} = IndexStyle(RA)
 
 DataAPI.refarray(sa::ScaledArray) = sa.refs
 DataAPI.refvalue(sa::ScaledArray, n::Integer) = getindex(DataAPI.refpool(sa), n)
 DataAPI.refpool(sa::ScaledArray) = sa.pool
+DataAPI.invrefpool(sa::ScaledArray) = sa.invpool
 
 DataAPI.refarray(ssa::SubArray{<:Any, <:Any, <:ScaledArray}) =
     view(parent(ssa).refs, ssa.indices...)
 DataAPI.refvalue(ssa::SubArray{<:Any, <:Any, <:ScaledArray}, n::Integer) =
     DataAPI.refvalue(parent(ssa), n)
 DataAPI.refpool(ssa::SubArray{<:Any, <:Any, <:ScaledArray}) =
     DataAPI.refpool(parent(ssa))
+DataAPI.invrefpool(ssa::SubArray{<:Any, <:Any, <:ScaledArray}) =
+    DataAPI.invrefpool(parent(ssa))
 
 @inline function Base.getindex(sa::ScaledArray, i::Int)
     refs = DataAPI.refarray(sa)
     @boundscheck checkbounds(refs, i)
     @inbounds n = refs[i]
+    iszero(n) && return missing
     pool = DataAPI.refpool(sa)
     @boundscheck checkbounds(pool, n)
     return @inbounds pool[n]
@@ -169,13 +207,14 @@ end
     @boundscheck checkbounds(refs, I...)
     @inbounds ns = refs[I...]
     pool = DataAPI.refpool(sa)
-    @boundscheck checkbounds(pool, ns)
+    N = length(pool)
+    @boundscheck checkindex(Bool, 0:N, ns) || throw_boundserror(pool, ns)
     return @inbounds pool[ns]
 end
 
 function Base.:(==)(x::ScaledArray, y::ScaledArray)
     size(x) == size(y) || return false
-    x.start == y.start && x.step == y.step && return x.refs == y.refs
+    first(x.pool) == first(y.pool) && step(x.pool) == step(y.pool) && return x.refs == y.refs
     eq = true
     for (p, q) in zip(x, y)
         # missing could arise

src/operations.jl

@@ -20,7 +20,17 @@ end
 
 # Obtain unique labels for row-wise pairs of values from a1 and a2 when mult is large enough
 function _mult!(a1::AbstractArray, a2::AbstractArray, mult)
-    a1 .+= mult .* (a2 .- 1)
+    z = zero(eltype(a1))
+    @inbounds for i in eachindex(a1)
+        x1 = a1[i]
+        x2 = a2[i]
+        # Handle missing values represented by zeros
+        if iszero(x1) || iszero(x2)
+            a1[i] = z
+        else
+            a1[i] += mult * (x2 - 1)
+        end
+	end
 end
 
 """
@@ -128,13 +138,24 @@ The returned array ensures well-defined time intervals for operations involving
 - `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)
-    eltype(time) <: ValidTimeType ||
-        throw(ArgumentError("unaccepted element type $(eltype(time)) from time column"))
-    step === nothing && (step = one(eltype(time)))
+    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)
     if rotation !== nothing
         refs = rotatingtime(rotation, time.refs)
-        time = ScaledArray(RefArray(refs), time.start, time.step, time.stop)
+        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
     return time
 end

src/procedures.jl

@@ -93,12 +93,16 @@ function overlap!(esample::BitVector, tr_rows::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))
     if !(timetype <: RotatingTimeValue)
         ecut = pr.ecut[1]
+        invpool === nothing || (ecut = invpool[ecut])
         valid_cohort = filter(x -> x < ecut || x in pr.e, overlap_time)
         filter!(x -> x < ecut, overlap_time)
     else
-        ecut = IdDict(e.rotation=>e.time for e in pr.ecut)
+        ecut = pr.ecut
+        invpool === nothing || (ecut = (invpool[e] for e in ecut))
+        ecut = IdDict(e.rotation=>e.time for e in ecut)
         valid_cohort = filter(x -> x.time < ecut[x.rotation] || x in pr.e, overlap_time)
         filter!(x -> x.time < ecut[x.rotation], overlap_time)
     end