Add MLE/MAP functionality

Project.toml

@@ -17,6 +17,7 @@ Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -41,6 +42,7 @@ ForwardDiff = "0.10.3"
 Libtask = "0.4"
 LogDensityProblems = "^0.9, 0.10"
 MCMCChains = "3.0.7"
+Optim = "0.20, 0.21"
 ProgressLogging = "0.1"
 Reexport = "0.2.0"
 Requires = "0.5, 1.0"
@@ -56,6 +58,8 @@ CmdStan = "593b3428-ca2f-500c-ae53-031589ec8ddd"
 DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73"
+NamedArrays = "86f7a689-2022-50b4-a561-43c23ac3c673"
+Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
@@ -66,4 +70,4 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Pkg", "PDMats", "TerminalLoggers", "Test", "UnicodePlots", "StatsBase", "FiniteDifferences", "DynamicHMC", "CmdStan", "BenchmarkTools", "Zygote", "ReverseDiff", "Memoization"]
+test = ["Pkg", "PDMats", "TerminalLoggers", "Test", "UnicodePlots", "StatsBase", "FiniteDifferences", "DynamicHMC", "CmdStan", "BenchmarkTools", "Zygote", "ReverseDiff", "Memoization", "Optim", "NamedArrays"]

src/Turing.jl

@@ -58,6 +58,11 @@ using .Variational
     end
 end
 
+@init @require Optim="429524aa-4258-5aef-a3af-852621145aeb" @eval begin
+    include("modes/ModeEstimation.jl")
+    export MAP, MLE, optimize
+end
+
 ###########
 # Exports #
 ###########
@@ -87,7 +92,7 @@ export  @model,                 # modelling
         CSMC,
         PG,
 
-        vi,                    # variational inference
+        vi,                     # variational inference
         ADVI,
 
         sample,                 # inference
@@ -110,5 +115,8 @@ export  @model,                 # modelling
         LogPoisson,
         NamedDist,
         filldist,
-        arraydist
+        arraydist,
+
+        MLE,                    # mode estimation tools
+        MAP
 end

src/inference/Inference.jl

@@ -305,19 +305,21 @@ Return a named tuple of parameters.
 getparams(t) = t.θ
 getparams(t::VarInfo) = tonamedtuple(TypedVarInfo(t))
 
-function _params_to_array(ts)
-    names_set = Set{String}()
+function _params_to_array(ts::Vector)
+    names = Vector{String}()
     # Extract the parameter names and values from each transition.
     dicts = map(ts) do t
         nms, vs = flatten_namedtuple(getparams(t))
         for nm in nms
-            push!(names_set, nm)
+            if !(nm in names)
+                push!(names, nm)
+            end
         end
         # Convert the names and values to a single dictionary.
         return Dict(nms[j] => vs[j] for j in 1:length(vs))
     end
-    names = collect(names_set)
-    vals = [get(dicts[i], key, missing) for i in eachindex(dicts),
+    # names = collect(names_set)
+    vals = [get(dicts[i], key, missing) for i in eachindex(dicts), 
         (j, key) in enumerate(names)]
 
     return names, vals

src/modes/ModeEstimation.jl

@@ -0,0 +1,284 @@
+using ..Turing
+using ..Bijectors
+using LinearAlgebra
+
+import ..AbstractMCMC: AbstractSampler
+import ..DynamicPPL
+import ..DynamicPPL: Model, AbstractContext, VarInfo, AbstractContext, VarName,
+    _getindex, getsym, getfield, settrans!,  setorder!,
+    get_and_set_val!, istrans, tilde, dot_tilde
+import Optim
+import Optim: optimize
+import NamedArrays
+import ..ForwardDiff
+import StatsBase
+import Printf
+
+struct MLE end
+struct MAP end
+
+"""
+    OptimizationContext{C<:AbstractContext} <: AbstractContext
+
+The `OptimizationContext` transforms variables to their constrained space, but
+does not use the density with respect to the transformation. This context is
+intended to allow an optimizer to sample in R^n freely.
+"""
+struct OptimizationContext{C<:AbstractContext} <: AbstractContext
+    context::C
+end
+
+# assume
+function DynamicPPL.tilde(ctx::OptimizationContext{<:LikelihoodContext}, spl, dist, vn::VarName, inds, vi)
+    r = vi[vn]
+    return r, 0
+end
+
+function DynamicPPL.tilde(ctx::OptimizationContext, spl, dist, vn::VarName, inds, vi)
+    r = vi[vn]
+    return r, Distributions.logpdf(dist, r)
+end
+
+# observe
+function DynamicPPL.tilde(ctx::OptimizationContext{<:PriorContext}, sampler, right, left, vi)
+    return 0
+end
+
+function DynamicPPL.tilde(ctx::OptimizationContext, sampler, dist, value, vi)
+    return Distributions.logpdf(dist, value)
+end
+
+# dot assume
+function DynamicPPL.dot_tilde(ctx::OptimizationContext{<:LikelihoodContext}, sampler, right, left, vn::VarName, _, vi)
+    vns, dist = get_vns_and_dist(right, left, vn)
+    r = getval(vi, vns)
+    return r, 0
+end
+
+function DynamicPPL.dot_tilde(ctx::OptimizationContext, sampler, right, left, vn::VarName, _, vi)
+    vns, dist = get_vns_and_dist(right, left, vn)
+    r = getval(vi, vns)
+    return r, loglikelihood(dist, r)
+end
+
+# dot observe
+function DynamicPPL.dot_tilde(ctx::OptimizationContext{<:PriorContext}, sampler, right, left, vn, _, vi)
+    return 0
+end
+
+function DynamicPPL.dot_tilde(ctx::OptimizationContext{<:PriorContext}, sampler, right, left, vi)
+    return 0
+end
+
+function DynamicPPL.dot_tilde(ctx::OptimizationContext, sampler, right, left, vn, _, vi)
+    vns, dist = get_vns_and_dist(right, left, vn)
+    r = getval(vi, vns)
+    return loglikelihood(dist, r)
+end
+
+function DynamicPPL.dot_tilde(ctx::OptimizationContext, sampler, dists, value, vi)
+    return sum(Distributions.logpdf.(dists, value))
+end
+
+function getval(
+    vi,
+    vns::AbstractVector{<:VarName},
+)
+    r = vi[vns]
+    return r
+end
+
+function getval(
+    vi,
+    vns::AbstractArray{<:VarName},
+)
+    r = reshape(vi[vec(vns)], size(vns))
+    return r
+end
+
+"""
+    OptimLogDensity{M<:Model,C<:Context,V<:VarInfo}
+
+A struct that stores the log density function of a `DynamicPPL` model.
+"""
+struct OptimLogDensity{M<:Model,C<:AbstractContext,V<:VarInfo}
+    "A `DynamicPPL.Model` constructed either with the `@model` macro or manually."
+    model::M
+    "A `DynamicPPL.AbstractContext` used to evaluate the model. `LikelihoodContext` or `DefaultContext` are typical for MAP/MLE."
+    context::C
+    "A `DynamicPPL.VarInfo` struct that will be used to update model parameters."
+    vi::V
+end
+
+"""
+    OptimLogDensity(model::Model, context::AbstractContext)
+
+Create a callable `OptimLogDensity` struct that evaluates a model using the given `context`.
+"""
+function OptimLogDensity(model::Model, context::AbstractContext)
+    init = VarInfo(model)
+    DynamicPPL.link!(init, DynamicPPL.SampleFromPrior())
+    return OptimLogDensity(model, context, init)
+end
+
+"""
+    (f::OptimLogDensity)(z)
+
+Evaluate the log joint (with `DefaultContext`) or log likelihood (with `LikelihoodContext`)
+at the array `z`.
+"""
+function (f::OptimLogDensity)(z)
+    spl = DynamicPPL.SampleFromPrior()
+
+    varinfo = DynamicPPL.VarInfo(f.vi, spl, z)
+    f.model(varinfo, spl, f.context)
+    return -DynamicPPL.getlogp(varinfo)
+end
+
+"""
+    ModeResult{
+        V<:NamedArrays.NamedArray, 
+        M<:NamedArrays.NamedArray, 
+        O<:Optim.MultivariateOptimizationResults, 
+        S<:NamedArrays.NamedArray
+    }
+
+A wrapper struct to store various results from a MAP or MLE estimation.
+"""
+struct ModeResult{
+    V<:NamedArrays.NamedArray, 
+    O<:Optim.MultivariateOptimizationResults,
+    M<:OptimLogDensity
+} <: StatsBase.StatisticalModel
+    "A vector with the resulting point estimates."
+    values :: V
+    "The stored Optim.jl results."
+    optim_result :: O
+    "The final log likelihood or log joint, depending on whether `MAP` or `MLE` was run."
+    lp :: Float64
+    "The evaluation function used to calculate the output."
+    f :: M
+end
+#############################
+# Various StatsBase methods #
+#############################
+
+
+
+function Base.show(io::IO, ::MIME"text/plain", m::ModeResult)
+    print(io, "ModeResult with minimized lp of ")
+    Printf.@printf(io, "%.2f", m.lp)
+    println(io)
+    show(io, m.values)
+end
+
+function Base.show(io::IO, m::ModeResult)
+    show(io, m.values.array)
+end
+
+function StatsBase.coeftable(m::ModeResult)
+    # Get columns for coeftable.
+    terms = StatsBase.coefnames(m)
+    estimates = m.values.array[:,1]
+    stderrors = StatsBase.stderror(m)
+    tstats = estimates ./ stderrors
+
+    StatsBase.CoefTable([estimates, stderrors, tstats], ["estimate", "stderror", "tstat"], terms)
+end
+
+function StatsBase.informationmatrix(m::ModeResult; hessian_function=ForwardDiff.hessian, kwargs...)
+    # Calculate Hessian and information matrix.
+    varnames = StatsBase.coefnames(m)
+    info = inv(hessian_function(m.f, m.values.array[:, 1]))
+    return NamedArrays.NamedArray(info, (varnames, varnames))
+end
+
+StatsBase.coef(m::ModeResult) = m.values
+StatsBase.coefnames(m::ModeResult) = names(m.values)[1]
+StatsBase.params(m::ModeResult) = StatsBase.coefnames(m)
+StatsBase.vcov(m::ModeResult) = StatsBase.informationmatrix(m)
+StatsBase.loglikelihood(m::ModeResult) = m.lp
+
+####################
+# Optim.jl methods #
+####################
+
+"""
+    Optim.optimize(model::Model, ::MLE, args...; kwargs...)
+
+Compute a maximum likelihood estimate of the `model`.
+
+# Examples
+
+```julia-repl
+@model function f(x)
+    m ~ Normal(0, 1)
+    x ~ Normal(m, 1)
+end
+
+model = f(1.5)
+mle = optimize(model, MLE())
+
+# Use a different optimizer
+mle = optimize(model, MLE(), NelderMead())
+```
+"""
+function Optim.optimize(model::Model, ::MLE, args...; kwargs...)
+    ctx = OptimizationContext(DynamicPPL.LikelihoodContext())
+    return optimize(model, OptimLogDensity(model, ctx), args...; kwargs...)
+end
+
+"""
+    Optim.optimize(model::Model, ::MAP, args...; kwargs...)
+
+Compute a maximum a posterior estimate of the `model`.
+
+# Examples
+
+```julia-repl
+@model function f(x)
+    m ~ Normal(0, 1)
+    x ~ Normal(m, 1)
+end
+
+model = f(1.5)
+map_est = optimize(model, MAP())
+
+# Use a different optimizer
+map_est = optimize(model, MAP(), NelderMead())
+```
+"""
+function Optim.optimize(model::Model, ::MAP, args...; kwargs...)
+    ctx = OptimizationContext(DynamicPPL.DefaultContext())
+    return optimize(model, OptimLogDensity(model, ctx), args...; kwargs...)
+end
+
+"""
+    Optim.optimize(model::Model, f::OptimLogDensity, optimizer=Optim.LBFGS(), args...; kwargs...)
+
+Estimate a mode, i.e., compute a MLE or MAP estimate.
+"""
+function Optim.optimize(model::Model, f::OptimLogDensity, optimizer=Optim.LBFGS(), args...; kwargs...)
+    # Do some initialization.
+    spl = DynamicPPL.SampleFromPrior()
+    init_vals = f.vi[spl]
+
+    # Optimize!
+    M = Optim.optimize(f, init_vals, optimizer, args...; kwargs...)
+
+    # Get the VarInfo at the MLE/MAP point, and run the model to ensure 
+    # correct dimensionality.
+    f.vi[spl] = M.minimizer
+    invlink!(f.vi, spl)
+    vals = f.vi[spl]
+    link!(f.vi, spl)
+
+    # Make one transition to get the parameter names.
+    ts = [Turing.Inference.Transition(DynamicPPL.tonamedtuple(f.vi), DynamicPPL.getlogp(f.vi))]
+    varnames, _ = Turing.Inference._params_to_array(ts)
+
+    # Store the parameters and their names in an array.
+    vmat = NamedArrays.NamedArray(vals, varnames)
+
+    return ModeResult(vmat, M, -M.minimum, f)
+end