mode_finding.jl

using Turing
using Optim
using UnicodePlots
import StatsBase.mode

@model gdemo_d() = begin
  s ~ InverseGamma(2, 3)
  m ~ Normal(0, sqrt(s))
  1.5 ~ Normal(m, sqrt(s))
  2.0 ~ Normal(m, sqrt(s))
  return s, m
end

function mode(model, contexts...)
  vi = DynamicPPL.VarInfo(model)
  sampler = DynamicPPL.SampleFromPrior()
  DynamicPPL.link!(vi, sampler)
  init = DynamicPPL.getall(vi)

  function f(p)
    varinfo = VarInfo(vi, sampler, p)
    return sum(c -> begin
        model(varinfo, sampler, c)
        -getlogp(varinfo)
        end, contexts)
  end
  
  mle = Optim.optimize(f, init, Optim.Options(iterations=1000, store_trace=true))
  DynamicPPL.setall!(vi, mle.minimizer)
  DynamicPPL.invlink!(vi, sampler)
  
  return DynamicPPL.getall(vi), map(v -> v.value, mle.trace), mle
end

model = gdemo_d()

# likelihood
mle, trace, m = mode(model, LikelihoodContext());
mle .- [0.0625031, 1.75] |> display
plt = lineplot(trace);

# prior
mle, trace, m = mode(model, PriorContext());
mle .- [mode(InverseGamma(2,3)), mode(Normal())] |> display
lineplot!(plt, trace);

# posterior
mle, trace, m = mode(model, PriorContext(), LikelihoodContext())