Allow sorting VecColumnTable and define TermSet as wrapped Set (#19)

Author: junyuan-chen

Project.toml

@@ -29,13 +29,11 @@ Reexport = "0.2, 1"
 StatsBase = "0.33"
 StatsModels = "0.6.18"
 Tables = "1.2"
-TypedTables = "1.2"
 julia = "1.3"
 
 [extras]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 
 [targets]
-test = ["DataFrames", "Test", "TypedTables"]
+test = ["DataFrames", "Test"]

src/DiffinDiffsBase.jl

@@ -10,13 +10,14 @@ using PooledArrays: _label
 using Reexport
 using StatsBase: Weights, uweights
 @reexport using StatsModels
+using StatsModels: Schema
 using Tables
 using Tables: AbstractColumns, table, istable, columnnames, getcolumn
 
 import Base: ==, show, union
 import Base: eltype, firstindex, lastindex, getindex, iterate, length, sym_in
 import StatsBase: coef, vcov, responsename, coefnames, weights, nobs, dof_residual
-import StatsModels: termvars
+import StatsModels: concrete_term, schema, termvars
 
 const TimeType = Int
 
@@ -27,6 +28,10 @@ export cb,
        ≊,
        exampledata,
 
+       VecColumnTable,
+       VecColsRow,
+       subcolumns,
+
        TreatmentSharpness,
        SharpDesign,
        sharp,
@@ -51,13 +56,11 @@ export cb,
        istreated,
 
        TermSet,
+       termset,
        eachterm,
        TreatmentTerm,
        treat,
 
-       VecColumnTable,
-       subcolumns,
-
        findcell,
        cellrows,
 
@@ -84,10 +87,10 @@ export cb,
        treatnames
 
 include("utils.jl")
+include("tables.jl")
 include("treatments.jl")
 include("parallels.jl")
 include("terms.jl")
-include("tables.jl")
 include("operations.jl")
 include("StatsProcedures.jl")
 include("procedures.jl")

src/did.jl

@@ -20,9 +20,8 @@ _key(::Any) = throw(ArgumentError("unacceptable positional arguments"))
 
 function _totermset!(args::Dict{Symbol,Any}, s::Symbol)
     if haskey(args, s) && !(args[s] isa TermSet)
-        ts = TermSet()
-        foreach(t->setindex!(ts, nothing, t), args[s])
-        args[s] = ts
+        arg = args[s]
+        args[s] = arg isa Symbol ? termset(arg) : termset(arg...)
     end
 end
 
@@ -185,7 +184,7 @@ All concrete subtypes of `DIDResult` are expected to have the following fields:
 - `yname::String`: name of the outcome variable (generated by `StatsModels.coefnames`).
 - `coefnames::Vector{String}`: names of all treatment coefficients and covariates.
 - `coefinds::Dict{String, Int}`: a map from `coefnames` to integer indices for retrieving estimates by name.
-- `treatinds::Table`: tabular descriptions of treatment coefficients in the order of `coefnames` (do not contain covariates).
+- `treatcells`: Tables.jl-compatible tabular descriptions of treatment coefficients in the order of `coefnames` (do not contain covariates).
 - `weights::Union{Symbol, Nothing}`: column name of the weight variable (if specified).
 """
 abstract type DIDResult <: StatisticalModel end
@@ -237,18 +236,36 @@ coef(r::DIDResult, i::Int) = r.coef[i]
 coef(r::DIDResult, inds) = [coef(r, ind) for ind in inds]
 
 """
+    coef(r::DIDResult, bys::Pair{Symbol,<:Function}...)
     coef(f::Function, r::DIDResult)
 
 Return a vector of point estimates for treatment coefficients
-selected based on whether `f` returns `true` or `false`
-for each corresponding row in `treatinds`.
+selected based on the specified functions that return either `true` or `false`.
+If `bys` are specified,
+each function is applied column-wise to each value of the specified column in `treatcells`.
+A coefficient is selected if the returned values from the functions are all `true`.
+If a single function `f` is provided without specifing column name,
+`f` is applied row-wise to each row of `treatcells`.
 
 !!! note
     This method only selects estimates for treatment coefficients.
     Covariates are not taken into account.
 """
+@inline function coef(r::DIDResult, @nospecialize(bys::Pair{Symbol,<:Function}...))
+    nby = length(bys)
+    by = bys[1]
+    inds = by[2].(getcolumn(r.treatcells, by[1]))
+    if nby > 1
+        for i in 2:nby
+            by = bys[i]
+            inds .&= by[2].(getcolumn(r.treatcells, by[1]))
+        end
+    end
+    return view(r.coef, 1:length(r.treatcells[1]))[inds]
+end
+
 @inline coef(f::Function, r::DIDResult) =
-    view(r.coef, 1:length(r.treatinds))[f.(r.treatinds)]
+    view(r.coef, 1:length(r.treatcells[1]))[f.(Tables.rows(r.treatcells))]
 
 """
     vcov(r::DIDResult)
@@ -282,19 +299,38 @@ function vcov(r::DIDResult, inds)
 end
 
 """
+    vcov(r::DIDResult, bys::Pair{Symbol,<:Function}...)
     vcov(f::Function, r::DIDResult)
 
 Return a variance-covariance matrix for treatment coefficients
-selected based on whether `f` returns `true` or `false`
-for each corresponding row in `treatinds`.
+selected based on the specified functions that return either `true` or `false`.
+If `bys` are specified,
+each function is applied column-wise to each value of the specified column in `treatcells`.
+A coefficient is selected if the returned values from the functions are all `true`.
+If a single function `f` is provided without specifing column name,
+`f` is applied row-wise to each row of `treatcells`.
 
 !!! note
     This method only selects estimates for treatment coefficients.
     Covariates are not taken into account.
 """
+@inline function vcov(r::DIDResult, @nospecialize(bys::Pair{Symbol,<:Function}...))
+    N = length(r.treatcells[1])
+    nby = length(bys)
+    by = bys[1]
+    inds = by[2].(getcolumn(r.treatcells, by[1]))
+    if nby > 1
+        for i in 2:nby
+            by = bys[i]
+            inds .&= by[2].(getcolumn(r.treatcells, by[1]))
+        end
+    end
+    return view(r.vcov, 1:N, 1:N)[inds, inds]
+end
+
 @inline function vcov(f::Function, r::DIDResult)
-    N = length(r.treatinds)
-    inds = f.(r.treatinds)
+    N = length(r.treatcells[1])
+    inds = f.(Tables.rows(r.treatcells))
     return view(r.vcov, 1:N, 1:N)[inds, inds]
 end
 
@@ -340,7 +376,7 @@ coefnames(r::DIDResult) = r.coefnames
 
 Return a vector of names for treatment coefficients.
 """
-treatnames(r::DIDResult) = r.coefnames[1:size(r.treatinds,1)]
+treatnames(r::DIDResult) = r.coefnames[1:size(r.treatcells,1)]
 
 """
     weights(r::DIDResult)

src/operations.jl

@@ -71,12 +71,12 @@ findcell(names, data, esample=Colon()) =
 A utility function for processing the object `refrows` returned by [`findcell`](@ref).
 Unique row values from `cols` corresponding to
 the keys in `refrows` are sorted lexicographically
-and stored as rows in a `Tables.MatrixTable`.
+and stored as rows in a new `VecColumnTable`.
 Groups of row indices from the values of `refrows` are permuted to
 match the order of row values and collected in a `Vector`.
 
 # Returns
-- `cells::MatrixTable`: unique row values from columns in `cols`.
+- `cells::VecColumnTable`: unique row values from columns in `cols`.
 - `rows::Vector{Vector{Int}}`: row indices for each combination.
 """
 function cellrows(cols::VecColumnTable, refrows::IdDict)
@@ -85,21 +85,26 @@ function cellrows(cols::VecColumnTable, refrows::IdDict)
     ncol = length(cols)
     ncell = length(refrows)
     rows = Vector{Vector{Int}}(undef, ncell)
-    cache = Matrix{Any}(undef, ncell, ncol+1)
+    columns = AbstractVector[Vector{eltype(c)}(undef, ncell) for c in cols]
+    refs = Vector{keytype(refrows)}(undef, ncell)
     r = 0
     @inbounds for (k, v) in refrows
         r += 1
-        cache[r, end] = k
         row1 = v[1]
+        refs[r] = k
         for c in 1:ncol
-            cache[r, c] = cols[c][row1]
+            columns[c][r] = cols[c][row1]
         end
     end
-    sorted = sortslices(cache, dims=1)
-    cells = table(sorted[:,1:ncol], header=columnnames(cols))
+    cells = VecColumnTable(columns, _names(cols), _lookup(cols))
+    p = sortperm(cells)
+    # Replace each column of cells with a new one in the sorted order
+    @inbounds for i in 1:ncol
+        columns[i] = cells[i][p]
+    end
     # Collect rows in the same order as cells
     @inbounds for i in 1:ncell
-        rows[i] = refrows[sorted[i,end]]
+        rows[i] = refrows[refs[p[i]]]
     end
     return cells, rows
 end

src/tables.jl

@@ -29,18 +29,33 @@ function VecColumnTable(columns::Vector{AbstractVector}, names::Vector{Symbol})
     return VecColumnTable(columns, names, lookup)
 end
 
-function VecColumnTable(data)
+function VecColumnTable(data, esample::Union{BitVector, Nothing}=nothing)
     Tables.istable(data) || throw(ArgumentError("input data is not Tables.jl-compatible"))
     names = collect(Tables.columnnames(data))
     ncol = length(names)
     columns = Vector{AbstractVector}(undef, ncol)
     cols = Tables.columns(data)
-    @inbounds for i in keys(names)
-        columns[i] = Tables.getcolumn(cols, i)
+    if esample === nothing
+        @inbounds for i in keys(names)
+            columns[i] = Tables.getcolumn(cols, i)
+        end
+    else
+        @inbounds for i in keys(names)
+            columns[i] = view(Tables.getcolumn(cols, i), esample)
+        end
     end
     return VecColumnTable(columns, names)
 end
 
+function VecColumnTable(cols::VecColumnTable, esample::Union{BitVector, Nothing}=nothing)
+    esample === nothing && return cols
+    columns = similar(_columns(cols))
+    @inbounds for i in keys(columns)
+        columns[i] = view(cols[i], esample)
+    end
+    return VecColumnTable(columns, _names(cols), _lookup(cols))
+end
+
 _columns(cols::VecColumnTable) = getfield(cols, :columns)
 _names(cols::VecColumnTable) = getfield(cols, :names)
 _lookup(cols::VecColumnTable) = getfield(cols, :lookup)
@@ -150,3 +165,58 @@ end
 
 subcolumns(data, names::Symbol, rows=Colon(); nomissing=true) =
     subcolumns(data, [names], rows, nomissing=nomissing)
+
+const VecColsRow = Tables.ColumnsRow{VecColumnTable}
+
+ncol(r::VecColsRow) = length(r)
+
+_rowhash(cols::Tuple{AbstractVector}, r::Int, h::UInt=zero(UInt))::UInt =
+    hash(cols[1][r], h)
+
+function _rowhash(cols::Tuple{Vararg{AbstractVector}}, r::Int, h::UInt=zero(UInt))::UInt
+    h = hash(cols[1][r], h)
+    _rowhash(Base.tail(cols), r, h)
+end
+
+# hash is implemented following DataFrames.DataFrameRow for getting unique row values
+# Column names are not taken into account
+Base.hash(r::VecColsRow, h::UInt=zero(UInt)) =
+    _rowhash(ntuple(c->Tables.getcolumns(r)[c], ncol(r)), Tables.getrow(r), h)
+
+# Column names are not taken into account
+function Base.isequal(r1::VecColsRow, r2::VecColsRow)
+    length(r1) == length(r2) || return false
+    return all(((a, b),) -> isequal(a, b), zip(r1, r2))
+end
+
+# Column names are not taken into account
+function Base.isless(r1::VecColsRow, r2::VecColsRow)
+    length(r1) == length(r2) ||
+        throw(ArgumentError("compared VecColsRow do not have the same length"))
+    for (a, b) in zip(r1, r2)
+        isequal(a, b) || return isless(a, b)
+    end
+    return false
+end
+
+Base.sortperm(cols::VecColumnTable; @nospecialize(kwargs...)) =
+    sortperm(collect(Tables.rows(cols)); kwargs...)
+
+# names and lookup are not copied
+function Base.sort(cols::VecColumnTable; @nospecialize(kwargs...))
+    p = sortperm(cols; kwargs...)
+    columns = similar(_columns(cols))
+    i = 0
+    @inbounds for col in cols
+        i += 1
+        columns[i] = col[p]
+    end
+    return VecColumnTable(columns, _names(cols), _lookup(cols))
+end
+
+function Base.sort!(cols::VecColumnTable; @nospecialize(kwargs...))
+    p = sortperm(cols; kwargs...)
+    @inbounds for col in cols
+        col .= col[p]
+    end
+end