plot.jl

using RecipesBase
using StatPlots

@userplot MeanPlot
@userplot AutocorPlot
@userplot HistogramPlot
@userplot DensityPlot
@userplot MixedDensityPlot
@userplot TracePlot

@recipe function f(p::MeanPlot)
  c, i = p.args
  seriestype := :line
  xaxis := "Iteration"
  yaxis := "Mean"
  MCMCChain.cummean(chain.value[:, i, :])
end

@recipe function f(p::AutocorPlot)
  c, i = p.args
  ac = MCMCChain.autocor(c, lags=collect(lags))
  xaxis := "Lag"
  yaxis := "Autocorrelation"
  ac.value[i, :, :]
end

@recipe function f(p::HistogramPlot)
  c, i = p.args
  seriestype := :histogram
  xaxis := "Value"
  yaxis := "Value"
  c.value[:, i, :]
end

@recipe function f(p::DensityPlot)
  c, i = p.args
  seriestype := :density
  xaxis := "Value"
  yaxis := "Density"
  c.value[:, i, :]
end

@recipe function f(p::TracePlot)
  c, i = p.args
  seriestype := :line
  xaxis := "Iteration"
  yaxis := "Density"
  c.value[:, i, :]
end

@recipe function f(p::MixedDensityPlot)
  c, i = p.args
  discrete = MCMCChain.indiscretesupport(c, barbounds)
  seriestype := discrete[i] ? :histogram : :density
  xaxis := "Value"
  yaxis := "Density"
  c.value[:, i, :]
end

#################### Posterior Plot Recipies ####################
@recipe function f(c::MCMCChain.AbstractChains,
                   ptypes;
                   maxlag = round(Int, 10 * log10(length(c.range))),
                   barbounds = (0, Inf))
  ptypes = ptypes isa Symbol ? [ptypes] : ptypes

  nrows, nvars, nchains = size(c.value)
  ntypes = length(ptypes)
  layout := (nvars, ntypes)
  indices = reshape(1:nvars*ntypes, ntypes, nvars)'
  lags = 0:maxlag

  legend --> false

  for (j, ptype) in enumerate(ptypes)
    for i in 1:nvars
      @series begin
          seriestype := ptype
          subplot := indices[i, j]
          title := c.names[i]
          labels := map(k -> "Chain $k", 1:nchains)
          c, i
      end
    end
  end
end