nicer plot

ext/MCMCChainsMakieExt.jl

@@ -4,75 +4,35 @@ import MCMCChains
 import Makie
 
 
-# @recipe(MCMCChains.Chains) do scene
-#     Theme()
-# end
-
-function Makie.plot!(chns::MCMCChains.Chains)
-    params = names(chns, :parameters)
-
-    n_chains = length(chains(chns))
-    n_samples = length(chns)
-
-    fig = Figure()
-
-    for (i, param) in enumerate(params)
-        ax = Axis(fig[i, 1]; ylabel=string(param))
-        for chain in 1:n_chains
-            values = chns[:, param, chain]
-            lines!(ax, 1:n_samples, values; label=string(chain))
-        end
-
-        hideydecorations!(ax; label=false)
-        if i < length(params)
-            hidexdecorations!(ax; grid=false)
-        else
-            ax.xlabel = "Iteration"
-        end
-    end
-
-    for (i, param) in enumerate(params)
-        ax = Axis(fig[i, 2]; ylabel=string(param))
-        for chain in 1:n_chains
-            values = chns[:, param, chain]
-            density!(ax, values; label=string(chain))
-        end
-
-        hideydecorations!(ax)
-        if i < length(params)
-            hidexdecorations!(ax; grid=false)
-        else
-            ax.xlabel = "Parameter estimate"
-        end
-    end
-
-    axes = [only(contents(fig[i, 2])) for i in 1:length(params)]
-    linkxaxes!(axes...)
-
-    axislegend(first(axes))
-
-    return fig
-end
-
-
 function MCMCChains.myplot(chns::T) where {T<:MCMCChains.Chains}
     params = names(chns, :parameters)
 
     n_chains = length(MCMCChains.chains(chns))
     n_samples = length(chns)
+    n_params = length(params)
 
-    colors = Makie.cgrad(:roma100, n_chains, categorical=true)
-    fig = Makie.Figure()
+    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)
+            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 < length(params)
+        if i < n_params
             Makie.hidexdecorations!(ax; grid=false)
         else
             ax.xlabel = "Iteration"
@@ -83,21 +43,31 @@ function MCMCChains.myplot(chns::T) where {T<:MCMCChains.Chains}
         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.7))
+            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 < length(params)
+        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:length(params)]
+    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)
+    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