Add plotting via Makie as Extension

Project.toml

@@ -26,6 +26,12 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 TableTraits = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
+[weakdeps]
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+
+[extensions]
+MCMCChainsMakieExt = "Makie"
+
 [compat]
 AbstractMCMC = "0.4, 0.5, 1.0, 2.0, 3.0, 4, 5"
 AxisArrays = "0.4.4"

ext/MCMCChainsMakieExt.jl

@@ -0,0 +1,74 @@
+module MCMCChainsMakieExt
+
+import MCMCChains
+import Makie
+
+
+function MCMCChains.trace(chns::T) where {T<:MCMCChains.Chains}
+    params = MCMCChains.names(chns, :parameters)
+
+    n_chains = length(MCMCChains.chains(chns))
+    n_samples = length(chns)
+    n_params = length(params)
+
+    colors = Makie.to_colormap(:tol_vibrant)
+    width = 600
+    height = max(400, 80 * n_params)
+
+    fig = Makie.Figure(; size=(width, height))
+
+    for (i, param) in enumerate(params)
+        ax = Makie.Axis(fig[i+1, 1]; ylabel=string(param))
+        for chain in 1:n_chains
+            values = chns[:, param, chain]
+            Makie.lines!(
+                ax,
+                1:n_samples,
+                values;
+                label=string(chain),
+                color=(colors[chain], 0.7),
+                linewidth=0.7
+            )
+        end
+
+        Makie.hideydecorations!(ax; label=false)
+        if i < n_params
+            Makie.hidexdecorations!(ax; grid=false)
+        else
+            ax.xlabel = "Iteration"
+        end
+    end
+
+    for (i, param) in enumerate(params)
+        ax = Makie.Axis(fig[i+1, 2]; ylabel=string(param))
+        for chain in 1:n_chains
+            values = chns[:, param, chain]
+            Makie.density!(
+                ax,
+                values;
+                label=string(chain),
+                color=(colors[chain], 0.1),
+                strokewidth=1,
+                strokecolor=(colors[chain], 0.7)
+            )
+        end
+
+        Makie.hideydecorations!(ax)
+        if i < n_params
+            Makie.hidexdecorations!(ax; grid=false)
+        else
+            ax.xlabel = "Parameter estimate"
+        end
+    end
+
+    axes = [only(Makie.contents(fig[i+1, 2])) for i in 1:n_params]
+    Makie.linkxaxes!(axes...)
+
+    Makie.Legend(fig[1, 1:2], first(axes), "Chain", orientation=:horizontal, titlehalign=:left, halign=:left, titleposition=:left)
+
+    Makie.rowgap!(fig.layout, 10)
+    Makie.colgap!(fig.layout, 10)
+
+    return fig
+end
+end

src/MCMCChains.jl

@@ -47,6 +47,7 @@ export rafterydiag
 export rstar
 
 export hpd
+export trace
 
 """
     Chains
@@ -85,4 +86,20 @@ include("plot.jl")
 include("tables.jl")
 include("rstar.jl")
 
+
+### trace function via Extensions still needs to define a base function
+function trace end
+
+function __init__()
+
+    Base.Experimental.register_error_hint(MethodError) do io, exc, argtypes, kwargs
+        if exc.f in [trace]
+            if isempty(methods(exc.f))
+                print(io, "\n$(exc.f) has no methods, yet. Makie has to be loaded for the plotting extension to be activated. Run `using Makie`, `using CairoMakie`, `using GLMakie` or any other package that also loads Makie.")
+            end
+        end
+    end
+end
+
+
 end # module