Store SummaryStats labels separately from data (#81)

* Refactor to store labels separate from data

* Set var_names to nothing by default

* Update tests

* Increment patch number

* Run downgrade CI on minimum Julia version

DimensionalData has a fix for latest Julia release, but not on its earliest supported version

* Update Project.toml

Author: sethaxen

.github/workflows/CI.yml

@@ -31,7 +31,7 @@ jobs:
         downgrade:
           - false
         include:
-          - version: '1'
+          - version: 'min'
             os: ubuntu-latest
             arch: x64
             downgrade: true

Project.toml

@@ -1,7 +1,7 @@
 name = "PosteriorStats"
 uuid = "7f36be82-ad55-44ba-a5c0-b8b5480d7aa5"
 authors = ["Seth Axen <seth@sethaxen.com> and contributors"]
-version = "0.4.2"
+version = "0.4.3-DEV"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"

src/summarize.jl

@@ -117,14 +117,16 @@ Base.@constprop :aggressive function summarize(
     stats_funs_and_names...;
     kind::Union{Symbol,Val}=:all,
     name::String="SummaryStats",
-    var_names=axes(data, 3),
+    var_names=nothing,
     kwargs...,
 )
-    length(var_names) == size(data, 3) || throw(
-        DimensionMismatch(
-            "length $(length(var_names)) of `var_names` does not match number of parameters $(size(data, 3)) in `data`.",
-        ),
-    )
+    var_names === nothing ||
+        length(var_names) == size(data, 3) ||
+        throw(
+            DimensionMismatch(
+                "length $(length(var_names)) of `var_names` does not match number of parameters $(size(data, 3)) in `data`.",
+            ),
+        )
     if isempty(stats_funs_and_names)
         return _summarize(data, default_summary_stats(kind; kwargs...), name, var_names)
     else

src/summarystats.jl

@@ -31,60 +31,64 @@ this will be used for the row labels or will be replaced with the `labels` if pr
     display. If not provided, then the column `$(LABEL_COLUMN_NAME)` in `data` will be
     used if it exists. Otherwise, the parameter names will be numeric indices.
 """
-struct SummaryStats{D,N<:AbstractString}
+struct SummaryStats{D,L<:Union{Nothing,AbstractVector},N<:AbstractString}
     data::D
+    labels::L
     name::N
-    function SummaryStats(data, name::N) where {N<:AbstractString}
-        _coltable = Tables.columntable(data)
-        # define default parameter names if not present, and set as first column
-        if !haskey(_coltable, LABEL_COLUMN_NAME)
-            data_cols = _coltable
-            labels = Base.OneTo(Tables.rowcount(data))
-        else
-            data_colnames = filter(k -> k !== LABEL_COLUMN_NAME, keys(_coltable))
-            data_cols = NamedTuple{data_colnames}(_coltable)
-            labels = _coltable[LABEL_COLUMN_NAME]
-        end
-        coltable = merge((; LABEL_COLUMN_NAME => labels), data_cols)
-        return new{typeof(coltable),N}(coltable, name)
-    end
 end
 
 function SummaryStats(
     data; labels::Union{AbstractVector,Nothing}=nothing, name::AbstractString="SummaryStats"
 )
+    _coltable = Tables.columntable(data)
     if labels !== nothing
         length(labels) == Tables.rowcount(data) || throw(
             DimensionMismatch(
                 "length $(length(labels)) of `labels` does not match number of rows $(Tables.rowcount(data)) in `data`.",
             ),
         )
-        data_with_varnames = merge(
-            Tables.columntable(data), (; LABEL_COLUMN_NAME => labels)
+    end
+    if haskey(_coltable, LABEL_COLUMN_NAME)
+        labels === nothing || throw(
+            ArgumentError(
+                "Either `labels` or a column named `$(LABEL_COLUMN_NAME)` may be provided, but not both.",
+            ),
         )
-    else
-        data_with_varnames = data
+        data_colnames = filter(k -> k !== LABEL_COLUMN_NAME, keys(_coltable))
+        data_cols = NamedTuple{data_colnames}(_coltable)
+        _labels = _coltable[LABEL_COLUMN_NAME]
+        return SummaryStats(data_cols, _labels, name)
     end
-    return SummaryStats(data_with_varnames, name)
+    return SummaryStats(_coltable, labels, name)
 end
 
 # forward key interfaces from its parent
 Base.parent(stats::SummaryStats) = getfield(stats, :data)
-Base.keys(stats::SummaryStats) = map(Symbol, Tables.columnnames(stats))
+Base.keys(stats::SummaryStats) = Tables.columnnames(stats)
 Base.haskey(stats::SummaryStats, nm::Symbol) = nm ∈ keys(stats)
-Base.length(stats::SummaryStats) = length(parent(stats))
+Base.length(stats::SummaryStats) = length(parent(stats)) + 1
 Base.getindex(stats::SummaryStats, i::Union{Int,Symbol}) = Tables.getcolumn(stats, i)
-Base.iterate(stats::SummaryStats, rest...) = iterate(parent(stats), rest...)
+Base.iterate(stats::SummaryStats) = (_labels(stats), 2)
+function Base.iterate(stats::SummaryStats, i::Int)
+    state = iterate(parent(stats), i - 1)
+    state === nothing && return nothing
+    return (state[1], state[2] + 1)
+end
 function Base.merge(stats::SummaryStats, other_stats::SummaryStats...)
     isempty(other_stats) && return stats
     stats_all = (stats, other_stats...)
     stats_last = last(stats_all)
-    return SummaryStats(merge(map(parent, stats_all)...), stats_last.name)
+    return SummaryStats(
+        merge(map(parent, stats_all)...),
+        getfield(stats_last, :labels),
+        getfield(stats_last, :name),
+    )
 end
 for f in (:(==), :isequal)
     @eval begin
         function Base.$(f)(stats::SummaryStats, other_stats::SummaryStats)
-            return $(f)(parent(stats), parent(other_stats))
+            return $(f)(_labels(stats), _labels(other_stats)) &&
+                   $(f)(parent(stats), parent(other_stats))
         end
     end
 end
@@ -99,30 +103,42 @@ function Base.show(io::IO, mime::MIME"text/html", stats::SummaryStats; kwargs...
 end
 
 function _show(io::IO, mime::MIME, stats::SummaryStats; kwargs...)
-    nt = parent(stats)
-    data = nt[keys(nt)[2:end]]
+    data = parent(stats)
     rhat_formatter = _prettytables_rhat_formatter(data)
     extra_formatters = rhat_formatter === nothing ? () : (rhat_formatter,)
     return _show_prettytable(
         io,
         mime,
         data;
         title=stats.name,
-        row_labels=Tables.getcolumn(stats, LABEL_COLUMN_NAME),
+        row_labels=_labels(stats),
         extra_formatters,
         kwargs...,
     )
 end
 
 #### Tables interface as column table
 
+_labels(s::SummaryStats) = getfield(s, :labels)
+_labels(s::SummaryStats{<:Any,Nothing}) = eachindex(values(parent(s))...)
+
 Tables.istable(::Type{<:SummaryStats}) = true
 Tables.columnaccess(::Type{<:SummaryStats}) = true
 Tables.columns(s::SummaryStats) = s
-Tables.columnnames(s::SummaryStats) = Tables.columnnames(parent(s))
-Tables.getcolumn(stats::SummaryStats, i::Int) = Tables.getcolumn(parent(stats), i)
-Tables.getcolumn(stats::SummaryStats, nm::Symbol) = Tables.getcolumn(parent(stats), nm)
-Tables.schema(s::SummaryStats) = Tables.schema(parent(s))
+Tables.columnnames(s::SummaryStats) = (LABEL_COLUMN_NAME, Tables.columnnames(parent(s))...)
+function Tables.getcolumn(stats::SummaryStats, i::Int)
+    i == 1 && return _labels(stats)
+    return Tables.getcolumn(parent(stats), i - 1)
+end
+function Tables.getcolumn(stats::SummaryStats, nm::Symbol)
+    nm === LABEL_COLUMN_NAME && return _labels(stats)
+    return Tables.getcolumn(parent(stats), nm)
+end
+function Tables.schema(s::SummaryStats)
+    labels = _labels(s)
+    sch = Tables.schema(parent(s))
+    return Tables.Schema((LABEL_COLUMN_NAME, sch.names...), (eltype(labels), sch.types...))
+end
 
 IteratorInterfaceExtensions.isiterable(::SummaryStats) = true
 function IteratorInterfaceExtensions.getiterator(s::SummaryStats)

test/summarystats.jl

@@ -13,35 +13,44 @@ using Test
 
     @testset "constructors" begin
         stats1 = SummaryStats(data)
-        @test stats1.data == data_with_default_labels
+        @test stats1.data == data
+        @test isnothing(stats1.labels)
         @test stats1.name == "SummaryStats"
 
         stats2 = SummaryStats(data; name="Stats")
-        @test stats2.data == data_with_default_labels
+        @test stats2.data == data
+        @test isnothing(stats2.labels)
         @test stats2.name == "Stats"
 
         stats3 = SummaryStats(data; labels)
-        @test stats3.data == data_with_labels
+        @test stats3.data == data
+        @test stats3.labels == labels
         @test stats3.name == "SummaryStats"
 
         stats4 = SummaryStats(data; labels, name="Stats")
-        @test stats4.data == data_with_labels
+        @test stats4.data == data
+        @test stats4.labels == labels
         @test stats4.name == "Stats"
 
         stats5 = SummaryStats(data_with_labels)
-        @test stats5.data == data_with_labels
+        @test stats5.data == data
+        @test stats5.labels == labels
 
         stats6 = SummaryStats(merge(data, (; label=labels)))
-        @test stats6.data == data_with_labels
+        @test stats6.data == data
+        @test stats6.labels == labels
         @test stats6.name == "SummaryStats"
+
+        @test_throws ArgumentError SummaryStats(data_with_labels; labels)
+        @test_throws DimensionMismatch SummaryStats(data; labels=labels[1:(end - 1)])
     end
 
     @inferred SummaryStats(data; name="Stats")
     stats_with_names(data, name, labels) = SummaryStats(data; name, labels)
     stats = @inferred stats_with_names(data, "Stats", labels)
 
     @testset "basic interfaces" begin
-        @test parent(stats) == data_with_labels
+        @test parent(stats) == data
         @test stats.name == "Stats"
         @test SummaryStats(data; name="MoreStats").name == "MoreStats"
         @test keys(stats) == (:label, keys(data)...)