Include parameter dictionary in optimisation return value (#2710)

Author: penelopeysm

HISTORY.md

@@ -1,3 +1,21 @@
+# 0.41.1
+
+The `ModeResult` struct returned by `maximum_a_posteriori` and `maximum_likelihood` can now be wrapped in `InitFromParams()`.
+This makes it easier to use the parameters in downstream code, e.g. when specifying initial parameters for MCMC sampling.
+For example:
+
+```julia
+@model function f()
+    # ...
+end
+model = f()
+opt_result = maximum_a_posteriori(model)
+
+sample(model, NUTS(), 1000; initial_params=InitFromParams(opt_result))
+```
+
+If you need to access the dictionary of parameters, it is stored in `opt_result.params` but note that this field may change in future breaking releases as that Turing's optimisation interface is slated for overhaul in the near future.
+
 # 0.41.0
 
 ## DynamicPPL 0.38

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.41.0"
+version = "0.41.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

ext/TuringOptimExt.jl

@@ -192,7 +192,7 @@ function _optimize(
     varnames = map(Symbol ∘ first, vns_vals_iter)
     vals = map(last, vns_vals_iter)
     vmat = NamedArrays.NamedArray(vals, varnames)
-    return Optimisation.ModeResult(vmat, M, -M.minimum, logdensity_optimum)
+    return Optimisation.ModeResult(vmat, M, -M.minimum, logdensity_optimum, vals_dict)
 end
 
 end # module

src/optimisation/Optimisation.jl

@@ -4,6 +4,7 @@ using ..Turing
 using NamedArrays: NamedArrays
 using AbstractPPL: AbstractPPL
 using DynamicPPL: DynamicPPL
+using DocStringExtensions: TYPEDFIELDS
 using LogDensityProblems: LogDensityProblems
 using Optimization: Optimization
 using OptimizationOptimJL: OptimizationOptimJL
@@ -154,13 +155,22 @@ end
         V<:NamedArrays.NamedArray,
         M<:NamedArrays.NamedArray,
         O<:Optim.MultivariateOptimizationResults,
-        S<:NamedArrays.NamedArray
+        S<:NamedArrays.NamedArray,
+        P<:AbstractDict{<:VarName,<:Any}
     }
 
 A wrapper struct to store various results from a MAP or MLE estimation.
+
+## Fields
+
+$(TYPEDFIELDS)
 """
-struct ModeResult{V<:NamedArrays.NamedArray,O<:Any,M<:OptimLogDensity} <:
-       StatsBase.StatisticalModel
+struct ModeResult{
+    V<:NamedArrays.NamedArray,
+    O<:Any,
+    M<:OptimLogDensity,
+    P<:AbstractDict{<:AbstractPPL.VarName,<:Any},
+} <: StatsBase.StatisticalModel
     "A vector with the resulting point estimates."
     values::V
     "The stored optimiser results."
@@ -169,6 +179,8 @@ struct ModeResult{V<:NamedArrays.NamedArray,O<:Any,M<:OptimLogDensity} <:
     lp::Float64
     "The evaluation function used to calculate the output."
     f::M
+    "Dictionary of parameter values"
+    params::P
 end
 
 function Base.show(io::IO, ::MIME"text/plain", m::ModeResult)
@@ -182,6 +194,15 @@ function Base.show(io::IO, m::ModeResult)
     return show(io, m.values.array)
 end
 
+"""
+    InitFromParams(m::ModeResult)
+
+Initialize a model from the parameters stored in a `ModeResult`.
+"""
+function DynamicPPL.InitFromParams(m::ModeResult)
+    return DynamicPPL.InitFromParams(m.params)
+end
+
 # Various StatsBase methods for ModeResult
 
 """
@@ -355,9 +376,13 @@ function ModeResult(log_density::OptimLogDensity, solution::SciMLBase.Optimizati
     iters = map(AbstractPPL.varname_and_value_leaves, keys(vals), values(vals))
     vns_vals_iter = mapreduce(collect, vcat, iters)
     syms = map(Symbol ∘ first, vns_vals_iter)
-    vals = map(last, vns_vals_iter)
+    split_vals = map(last, vns_vals_iter)
     return ModeResult(
-        NamedArrays.NamedArray(vals, syms), solution, -solution.objective, log_density
+        NamedArrays.NamedArray(split_vals, syms),
+        solution,
+        -solution.objective,
+        log_density,
+        vals,
     )
 end
 

test/optimisation/Optimisation.jl

@@ -101,6 +101,13 @@ using Turing
                 @test result.optim_result.retcode == Optimization.ReturnCode.Success
             end
             @test isapprox(result.lp, true_logp, atol=0.01)
+            # check that the parameter dict matches the NamedArray
+            # NOTE: This test only works for models where all parameters are identity
+            # varnames AND real-valued. Thankfully, this is true for `gdemo`.
+            @test length(only(result.values.dicts)) == length(result.params)
+            for (k, index) in only(result.values.dicts)
+                @test result.params[AbstractPPL.VarName{k}()] == result.values.array[index]
+            end
         end
 
         @testset "MLE" begin
@@ -546,6 +553,26 @@ using Turing
         end
     end
 
+    @testset "using ModeResult to initialise MCMC" begin
+        @model function f(y)
+            μ ~ Normal(0, 1)
+            σ ~ Gamma(2, 1)
+            return y ~ Normal(μ, σ)
+        end
+        model = f(randn(10))
+        mle = maximum_likelihood(model)
+        # TODO(penelopeysm): This relies on the fact that HMC does indeed
+        # use the initial_params passed to it. We should use something
+        # like a StaticSampler (see test/mcmc/Inference) to make this more
+        # robust.
+        chain = sample(
+            model, HMC(0.1, 10), 2; initial_params=InitFromParams(mle), num_warmup=0
+        )
+        # Check that those parameters were indeed used as initial params
+        @test chain[:µ][1] == mle.params[@varname(µ)]
+        @test chain[:σ][1] == mle.params[@varname(σ)]
+    end
+
     # Issue: https://discourse.julialang.org/t/turing-mixture-models-with-dirichlet-weightings/112910
     @testset "Optimization with different linked dimensionality" begin
         @model demo_dirichlet() = x ~ Dirichlet(2 * ones(3))
@@ -621,7 +648,13 @@ using Turing
         m = saddle_model()
         optim_ld = Turing.Optimisation.OptimLogDensity(m, DynamicPPL.getloglikelihood)
         vals = Turing.Optimisation.NamedArrays.NamedArray([0.0, 0.0])
-        m = Turing.Optimisation.ModeResult(vals, nothing, 0.0, optim_ld)
+        m = Turing.Optimisation.ModeResult(
+            vals,
+            nothing,
+            0.0,
+            optim_ld,
+            Dict{AbstractPPL.VarName,Float64}(@varname(x) => 0.0, @varname(y) => 0.0),
+        )
         ct = coeftable(m)
         @assert isnan(ct.cols[2][1])
         @assert ct.colnms[end] == "Error notes"