Import and extend PosteriorStats

Project.toml

@@ -8,6 +8,7 @@ version = "7.3.0"
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -18,7 +19,7 @@ MCMCDiagnosticTools = "be115224-59cd-429b-ad48-344e309966f0"
 MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 NaturalSort = "c020b1a1-e9b0-503a-9c33-f039bfc54a85"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
-PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
+PosteriorStats = "7f36be82-ad55-44ba-a5c0-b8b5480d7aa5"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -30,6 +31,7 @@ Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 [compat]
 AbstractMCMC = "0.4, 0.5, 1.0, 2.0, 3.0, 4, 5"
 AxisArrays = "0.4.4"
+Compat = "4.2.0"
 DataAPI = "1.16.0"
 Dates = "<0.0.1, 1"
 Distributions = "0.21, 0.22, 0.23, 0.24, 0.25"
@@ -40,12 +42,12 @@ MCMCDiagnosticTools = "0.3"
 MLJModelInterface = "0.3.5, 0.4, 1.0"
 NaturalSort = "1"
 OrderedCollections = "1.4"
-PrettyTables = "0.9, 0.10, 0.11, 0.12, 1, 2"
+PosteriorStats = "0.3"
 Random = "<0.0.1, 1"
 RecipesBase = "0.7, 0.8, 1.0"
 Statistics = "<0.0.1, 1"
 StatsBase = "0.33.2, 0.34"
 StatsFuns = "0.8, 0.9, 1"
 TableTraits = "0.4, 1"
 Tables = "1"
-julia = "1.6.3"
+julia = "1.10"

docs/Project.toml

@@ -3,20 +3,24 @@ CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656"
 Gadfly = "c91e804a-d5a3-530f-b6f0-dfbca275c004"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 MLJDecisionTreeInterface = "c6f25543-311c-4c74-83dc-3ea6d1015661"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd"
 
 [compat]
 CairoMakie = "0.6 - 0.13.0, 0.13.2 - 0.15"
 CategoricalArrays = "0.8, 0.9, 0.10, 1"
 DataFrames = "0.22, 1"
 Documenter = "0.26, 0.27, 1"
+DocumenterInterLinks = "1"
 Gadfly = "1.3.4"
 MCMCChains = "7"
 MLJBase = "0.19, 0.20, 0.21, 1"
 MLJDecisionTreeInterface = "0.3, 0.4"
+StableRNGs = "1"
 StatsPlots = "0.14, 0.15"
 julia = "1.7"

docs/make.jl

@@ -1,4 +1,5 @@
 using Documenter
+using DocumenterInterLinks
 using MCMCChains
 
 DocMeta.setdocmeta!(
@@ -8,6 +9,8 @@ DocMeta.setdocmeta!(
     recursive=true
 )
 
+links = InterLinks("PosteriorStats" => "https://julia.arviz.org/PosteriorStats/stable/")
+
 makedocs(
     sitename = "MCMCChains.jl",
     pages = [
@@ -28,5 +31,6 @@ makedocs(
     ],
     format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
     modules = [MCMCChains],
+    plugins = [links],
     checkdocs = :exports,
 )

docs/src/stats.md

@@ -8,5 +8,6 @@ describe
 mean
 summarystats
 quantile
-hpd
+eti
+hdi
 ```

docs/src/statsplots.md

@@ -174,13 +174,13 @@ ridgelineplot(chn, [:C, :B, :A])
 ## Forest
 
 ```@example statsplots
-forestplot(chn, [:C, :B, :A], hpd_val = [0.05, 0.15, 0.25])
+forestplot(chn, [:C, :B, :A], ci_probs = [0.05, 0.15, 0.25])
 ```
 
 ## Caterpillar
 
 ```@example statsplots
-forestplot(chn, chn.name_map[:parameters], hpd_val = [0.05, 0.15, 0.25], ordered = true)
+forestplot(chn, chn.name_map[:parameters], ci_fun=hdi, ci_probs = [0.05, 0.15, 0.25], ordered = true)
 ```
 
 ## API

docs/src/summarize.md

@@ -2,7 +2,6 @@
 
 The methods listed below are defined in `src/summarize.jl`.
 
-```@autodocs
-Modules = [MCMCChains]
-Pages = ["summarize.jl"]
+```@docs
+summarize
 ```

src/MCMCChains.jl

@@ -1,5 +1,6 @@
 module MCMCChains
 
+using Compat: stack
 using AxisArrays
 const axes = Base.axes
 import AbstractMCMC
@@ -25,7 +26,7 @@ import MCMCDiagnosticTools
 import MLJModelInterface
 import NaturalSort
 import OrderedCollections
-import PrettyTables
+import PosteriorStats
 import StatsFuns
 import Tables
 import TableTraits
@@ -40,8 +41,6 @@ export setrange, resetrange
 export set_section, get_params, sections, sort_sections, setinfo
 export replacenames, namesingroup, group
 export autocor, describe, sample, summarystats, AbstractWeights, mean, quantile
-export ChainDataFrame
-export summarize
 
 # Reexport diagnostics functions
 using MCMCDiagnosticTools:
@@ -68,7 +67,9 @@ export mcse
 export rafterydiag
 export rstar
 
-export hpd
+# Reexport stats functions
+using PosteriorStats: SummaryStats, eti, hdi, summarize
+export SummaryStats, eti, hdi, summarize
 
 """
     Chains

src/chains.jl

@@ -345,7 +345,7 @@ function Base.show(io::IO, chains::Chains)
 end
 
 function Base.show(io::IO, mime::MIME"text/plain", chains::Chains)
-    print(io, "Chains ", chains, ":\n\n", header(chains))
+    println(io, "Chains ", chains, ":\n\n", header(chains))
     println(io, "\nUse `describe(chains)` for summary statistics and quantiles.")
 end
 

src/discretediag.jl

@@ -17,16 +17,22 @@ function MCMCDiagnosticTools.discretediag(
         _permutedims_diagnostics(_chains.value.data); kwargs...
     )
 
-    # Create dataframes
-    parameters = (parameters = names(_chains),)
-    between_chain_df = ChainDataFrame(
-        "Chisq diagnostic - Between chains", merge(parameters, between_chain_vals),
+    # Create SummaryStats
+    param_names = names(_chains)
+    between_chain_stats = SummaryStats(
+        between_chain_vals;
+        name = "Chisq diagnostic - Between chains",
+        labels = param_names,
     )
-    within_chain_dfs = map(1:size(_chains, 3)) do i
+    within_chain_stats = map(1:size(_chains, 3)) do i
         vals = map(val -> val[:, i], within_chain_vals)
-        return ChainDataFrame("Chisq diagnostic - Chain $i", merge(parameters, vals))
+        return SummaryStats(
+            vals;
+            name = "Chisq diagnostic - Chain $i",
+            labels = param_names,
+        )
     end
-    dfs = vcat(between_chain_df, within_chain_dfs)
+    stats = vcat([between_chain_stats], within_chain_stats)
 
-    return dfs
+    return stats
 end

src/ess_rhat.jl

@@ -24,9 +24,9 @@ function MCMCDiagnosticTools.ess(
 
     # Convert to NamedTuple
     ess_per_sec = ess ./ dur
-    nt = merge((parameters = names(_chains),), (; ess, ess_per_sec))
+    nt = (; ess, ess_per_sec)
 
-    return ChainDataFrame("ESS", nt)
+    return SummaryStats(nt; name = "ESS", labels = names(_chains))
 end
 
 """
@@ -48,9 +48,9 @@ function MCMCDiagnosticTools.rhat(
     )
 
     # Convert to NamedTuple
-    nt = merge((parameters = names(_chains),), (; rhat))
+    nt = (; rhat)
 
-    return ChainDataFrame("R-hat", nt)
+    return SummaryStats(nt; name = "R-hat", labels = names(_chains))
 end
 
 """
@@ -79,7 +79,7 @@ function MCMCDiagnosticTools.ess_rhat(
 
     # Convert to NamedTuple
     ess_per_sec = ess_rhat.ess ./ dur
-    nt = merge((parameters = names(_chains),), ess_rhat, (; ess_per_sec))
+    nt = merge(ess_rhat, (; ess_per_sec))
 
-    return ChainDataFrame("ESS/R-hat", nt)
+    return SummaryStats(nt; name = "ESS/R-hat", labels = names(_chains))
 end

src/gelmandiag.jl

@@ -12,12 +12,13 @@ function MCMCDiagnosticTools.gelmandiag(
     results = MCMCDiagnosticTools.gelmandiag(_permutedims_diagnostics(psi); kwargs...)
 
     # Create a data frame with the results.
-    df = ChainDataFrame(
-        "Gelman, Rubin, and Brooks diagnostic",
-        merge((parameters = names(_chains),), results),
+    stats = SummaryStats(
+        results;
+        name = "Gelman, Rubin, and Brooks diagnostic",
+        labels = names(_chains),
     )
 
-    return df
+    return stats
 end
 
 function MCMCDiagnosticTools.gelmandiag_multivariate(
@@ -36,11 +37,12 @@ function MCMCDiagnosticTools.gelmandiag_multivariate(
         kwargs...,
     )
 
-    # Create a data frame with the results.
-    df = ChainDataFrame(
-        "Gelman, Rubin, and Brooks diagnostic",
-        (parameters = names(_chains), psrf = results.psrf, psrfci = results.psrfci),
+    # Create SummaryStats with the results.
+    stats = SummaryStats(
+        (psrf = results.psrf, psrfci = results.psrfci);
+        name = "Gelman, Rubin, and Brooks diagnostic",
+        labels = names(_chains),
     )
 
-    return df, results.psrfmultivariate
+    return stats, results.psrfmultivariate
 end

src/gewekediag.jl

@@ -18,12 +18,14 @@ function MCMCDiagnosticTools.gewekediag(
         return namedtuple_of_vecs
     end
 
-    # Create data frames.
-    parameters = (parameters = names(_chains),)
-    dfs = [
-        ChainDataFrame("Geweke diagnostic - Chain $i", merge(parameters, result))
-        for (i, result) in enumerate(results)
+    # Create SummaryStats.
+    stats = [
+        SummaryStats(
+            result;
+            name = "Geweke diagnostic - Chain $i",
+            labels = names(_chains),
+        ) for (i, result) in enumerate(results)
     ]
 
-    return dfs
+    return stats
 end

src/heideldiag.jl

@@ -16,14 +16,14 @@ function MCMCDiagnosticTools.heideldiag(
         return namedtuple_of_vecs
     end
 
-    # Create data frames.
-    parameters = (parameters = names(_chains),)
-    dfs = [
-        ChainDataFrame(
-            "Heidelberger and Welch diagnostic - Chain $i", merge(parameters, result)
-        )
-        for (i, result) in enumerate(results)
+    # Create SummaryStats.
+    stats = [
+        SummaryStats(
+            result;
+            name = "Heidelberger and Welch diagnostic - Chain $i",
+            labels = names(_chains),
+        ) for (i, result) in enumerate(results)
     ]
 
-    return dfs
+    return stats
 end

src/mcse.jl

@@ -1,5 +1,5 @@
 """
-    mcse(chains::Chains; duration=compute_duration, kwargs...)
+    mcse(chains::Chains; kwargs...)
 
 Estimate the Monte Carlo standard error.
 """
@@ -16,7 +16,7 @@ function MCMCDiagnosticTools.mcse(
         kwargs...,
     )
 
-    nt = merge((parameters = names(_chains),), (; mcse))
+    nt = (; mcse)
 
-    return ChainDataFrame("MCSE", nt)
+    return SummaryStats(nt; name = "MCSE", labels = names(_chains))
 end