Prior/posterior predictive check plots

Project.toml

@@ -3,7 +3,7 @@ uuid = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 keywords = ["markov chain monte carlo", "probablistic programming"]
 license = "MIT"
 desc = "Chain types and utility functions for MCMC simulations."
-version = "4.17.0"
+version = "4.16.0"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/MCMCChains.jl

@@ -1,5 +1,6 @@
 module MCMCChains
 
+using Base: cumulative_compile_time_ns_after
 using AxisArrays
 const axes = Base.axes
 using AbstractFFTs
@@ -13,7 +14,7 @@ using Formatting
 using Dates
 using KernelDensity: kde, pdf
 import StatsBase: autocov, counts, sem, AbstractWeights,
-    autocor, describe, quantile, sample, summarystats, cov
+    autocor, describe, quantile, sample, summarystats, cov, ecdf
 import MLJModelInterface
 import NaturalSort
 import PrettyTables

src/plot.jl

@@ -6,12 +6,14 @@
 @shorthands corner
 @userplot RidgelinePlot
 @userplot ForestPlot
+@shorthands ppcplot
 
 struct _TracePlot; c; val; end
 struct _MeanPlot; c; val;  end
 struct _DensityPlot; c; val;  end
 struct _HistogramPlot; c; val;  end
 struct _AutocorPlot; lags; val;  end
+struct _PPCPlot; y_obs; y_pred; ymean_pred; end
 
 # define alias functions for old syntax
 const translationdict = Dict(
@@ -20,7 +22,8 @@ const translationdict = Dict(
                         :density => _DensityPlot,
                         :histogram => _HistogramPlot,
                         :autocorplot => _AutocorPlot,
-                        :pooleddensity => _DensityPlot
+                        :pooleddensity => _DensityPlot,
+                        :ppcplot => _PPCPlot
                       )
 
 const supportedplots = push!(collect(keys(translationdict)), :mixeddensity, :corner)
@@ -397,3 +400,178 @@ end
         end
     end
 end
+
+@recipe function f(
+    yobs_data,
+    ypred_data::Chains;
+    yvar_name::AbstractVector{Symbol} = [],
+    plot_type = :density,
+    predictive_check = :posterior,
+    n_samples::Int = 50
+    )
+
+    st = get(plotattributes, :seriestype, :traceplot)
+
+    if st == :ppcplot
+        N = n_samples <= size(ypred_data)[1] ? n_samples : size(ypred_data)[1]
+        index = sample(1:size(ypred_data)[1], N, replace = false, ordered = true)
+
+        if ndims(yobs_data) == 1
+            y_obs = plot_type == :cumulative ? ecdf(vec(yobs_data)) : vec(yobs_data)
+            predictions = ypred_data.value.data[index,:,:]
+            ymean_pred = (plot_type == :cumulative
+                ? ecdf(vec(mean(ypred_data.value.data, dims = 1)))
+                : vec(mean(ypred_data.value.data, dims = 1)))
+
+            if plot_type == :density || plot_type == :cumulative
+                if predictive_check == :posterior
+                    title --> "Posterior predictive check"
+                elseif predictive_check == :prior
+                    title --> "Prior predictive check"
+                else
+                    throw(ArgumentError("`predictive_check` must be one of `prior` or `posterior`"))
+                end
+                for i in 1:N
+                    y_pred = (plot_type == :cumulative ? ecdf(vec(predictions[i,:,:]))
+                        : vec(predictions[i,:,:]))
+                    ypred_label = (isempty(yvar_name) ? (i == 1 ? "y pred" : nothing)
+                        : (i == 1 ? "$(yvar_name[1]) pred" : nothing))
+                    @series begin
+                        seriestype := :density
+                        seriesalpha --> 0.3
+                        linecolor --> "#BBBBBB"
+                        label --> ypred_label
+                        y_pred
+                    end
+                end
+                @series begin
+                    seriestype := :density
+                    label --> (isempty(yvar_name) ? "y obs" : "$(yvar_name[1]) obs")
+                    y_obs
+                end
+                @series begin
+                    seriestype := :density
+                    label --> (isempty(yvar_name) ? "y mean" : "$(yvar_name[1]) mean")
+                    ymean_pred
+                end
+
+            elseif plot_type == :histogram
+                layout --> N + 2
+                k = 1
+                @series begin
+                    subplot := k
+                    seriestype := :histogram
+                    label --> (isempty(yvar_name) ? "y obs" : "$(yvar_name[1]) obs")
+                    y_obs
+                end
+                k = 2
+                @series begin
+                    subplot := k
+                    seriestype := :histogram
+                    label --> (isempty(yvar_name) ? "y mean" : "$(yvar_name[1]) mean")
+                    ymean_pred
+                end
+                for i in 1:N
+                    y_pred = predictions[i,:,:]
+                    @series begin
+                        subplot := k + i
+                        seriestype := :histogram
+                        label --> nothing
+                        y_pred
+                    end
+                end
+            else
+                throw(ArgumentError("`plot_type` must be one of `:density`, `:cumulative` or `histogram`"))
+            end
+
+        elseif ndims(yobs_data) > 1
+            n_yval = size(yobs_data)[1]
+            n_yvar = size(yobs_data)[2]
+            mean_arr = reshape(mean(ypred_data.value.data, dims = 1), (n_yval, n_yvar))
+            k = 0
+            for j in 1:n_yvar
+                sections = MCMCChains.group(ypred_data, Symbol(yvar_name[j]))
+                predictions = sections.value.data[index,:,:]
+                y_obs = (plot_type == :cumulative ? ecdf(vec(yobs_data[:,j]))
+                    : vec(yobs_data[:,j]))
+                ymean_pred = (plot_type == :cumulative ? ecdf(vec(mean_arr[:,j]))
+                    : vec(mean_arr[:,j]))
+
+                if plot_type == :density || plot_type == :cumulative
+                    k += 1
+                    layout --> (1, n_yvar)
+                    if predictive_check == :posterior
+                        title --> "Posterior predictive check"
+                    elseif predictive_check == :prior
+                        title --> "Prior predictive check"
+                    else
+                        throw(ArgumentError("`predictive_check` must be one of `prior` or `posterior`"))
+                    end
+
+                    for i in 1:N
+                        y_pred = (plot_type == :cumulative ? ecdf(vec(predictions[i,:,:]))
+                            : vec(predictions[i,:,:]))
+                        @series begin
+                            subplot := k
+                            seriestype := :density
+                            seriesalpha --> 0.3
+                            linecolor --> "#BBBBBB"
+                            label --> (i == 1 ? "$(yvar_name[j]) pred" : nothing)
+                            y_pred
+                        end
+                    end
+                    @series begin
+                        subplot := k
+                        seriestype := :density
+                        label --> "$(yvar_name[j]) obs"
+                        y_obs
+                    end
+                    @series begin
+                        subplot := k
+                        seriestype := :density
+                        label --> "$(yvar_name[j]) mean"
+                        ymean_pred
+                    end
+
+                elseif plot_type == :histogram
+                    subplot := k
+                    layout --> N + 2
+                    h = 1
+                    @series begin
+                        subplot := h
+                        seriestype := :histogram
+                        label --> "$(yvar_name[j]) obs"
+                        y_obs
+                    end
+                    h = 2
+                    @series begin
+                        subplot := h
+                        seriestype := :histogram
+                        label --> "$(yvar_name[j]) mean"
+                        ymean_pred
+                    end
+                    for i in 1:N
+                        y_pred = predictions[i,:,:]
+                        @series begin
+                            subplot := h + i
+                            seriestype := :histogram
+                            label --> nothing
+                            y_pred
+                        end
+                    end
+
+                else
+                    throw(ArgumentError("`plot_type` must be one of `:density`, `:cumulative` or `:histogram`"))
+                end
+            end
+        else
+            throw(ArgumentError("Observed data must have `dim > 1`"))
+        end
+    else
+
+    end
+end
+
+@recipe function f(p::_PPCPlot)
+    p.y_obs, p.y_pred
+end